1import torch
2import sentencepiece as spm
3from huggingface_hub import hf_hub_download
4from transformers import AutoModelForCausalLM, AutoConfig
56# Load model7model = AutoModelForCausalLM.from_pretrained(8"Yaongi/HybriKo-117M-ToolLLaMA-SFT",9 trust_remote_code=True,10 torch_dtype=torch.float32
11)12device ="cuda"if torch.cuda.is_available()else"cpu"13model.to(device)14model.eval()1516# Load tokenizer17sp_path = hf_hub_download("Yaongi/HybriKo-117M-ToolLLaMA-SFT","HybriKo_tok.model")18sp = spm.SentencePieceProcessor()19sp.Load(sp_path)2021# Special tokens22SPECIAL_TOKENS ={23"<|im_start|>":32000,24"<|im_end|>":32001,25"<thought>":32002,26"</thought>":32003,27"<tool_call>":32004,28"</tool_call>":32005,29"<tools>":32006,30"</tools>":32007,31}3233defencode(text):34"""Encode text with special token handling."""35for token, token_id in SPECIAL_TOKENS.items():36 text = text.replace(token,f" \x00{token_id}\x00 ")3738 tokens =[]39for part in text.split("\x00"):40if part.strip().isdigit()andint(part.strip())in SPECIAL_TOKENS.values():41 tokens.append(int(part.strip()))42elif part.strip():43 tokens.extend(sp.EncodeAsIds(part))44return tokens
4546defdecode(ids):47"""Decode token IDs to text."""48 id_to_token ={v: k for k, v in SPECIAL_TOKENS.items()}49 result =[]50 regular_ids =[]5152foridin ids:53ifidin id_to_token:54if regular_ids:55 result.append(sp.DecodeIds(regular_ids))56 regular_ids =[]57 result.append(id_to_token[id])58else:59 regular_ids.append(id)6061if regular_ids:62 result.append(sp.DecodeIds(regular_ids))6364return"".join(result)6566@torch.no_grad()67defgenerate(prompt, max_new_tokens=200, temperature=0.7, top_k=50):68"""Generate with stop sequence detection."""69 input_ids = torch.tensor([encode(prompt)]).to(device)70 stop_sequences =["<|im_end|>","</tool_call>"]7172for _ inrange(max_new_tokens):73 logits = model(input_ids)["logits"][:,-1]/ temperature
7475if top_k:76 v, _ = torch.topk(logits,min(top_k, logits.size(-1)))77 logits[logits < v[:,[-1]]]=float("-inf")7879 probs = torch.softmax(logits, dim=-1)80 next_token = torch.multinomial(probs,1)81 input_ids = torch.cat([input_ids, next_token], dim=1)8283# Check stop sequences84 text = decode(input_ids[0].tolist())85for stop in stop_sequences:86if stop in text.split(prompt)[-1]:87return text
8889return decode(input_ids[0].tolist())
Example: Weather API Call
python
1prompt ="""<|im_start|>system
2You are an AI assistant with access to tools.
3<tools>
4{"name": "get_weather", "description": "Get current weather", "parameters": {"location": {"type": "string", "description": "City name"}}}
5</tools><|im_end|>
6<|im_start|>user
7What's the weather in Seoul?<|im_end|>
8<|im_start|>assistant
9"""1011response = generate(prompt, temperature=0.3, top_k=10)12print(response)
Expected Output:
<|im_start|>assistant
<thought>
The user wants to know the weather in Seoul. I should call the get_weather function.
</thought>
<tool_call>
{"name": "get_weather", "arguments": {"location": "Seoul"}}
</tool_call><|im_end|>
Example: Search Query
python
1prompt ="""<|im_start|>system
2You are an AI assistant with access to tools.
3<tools>
4{"name": "web_search", "description": "Search the web", "parameters": {"query": {"type": "string"}}}
5</tools><|im_end|>
6<|im_start|>user
7Find information about the latest AI research<|im_end|>
8<|im_start|>assistant
9"""1011response = generate(prompt, temperature=0.3)12print(response)
Prompt Format (Hermes ChatML)
<|im_start|>system
You are an AI assistant with access to tools.
<tools>
[Tool definitions in JSON format]
</tools><|im_end|>
<|im_start|>user
[User message]<|im_end|>
<|im_start|>assistant
<thought>
[Model's reasoning]
</thought>
<tool_call>
{"name": "[tool_name]", "arguments": {...}}
</tool_call><|im_end|}
<|im_start|>tool
[Tool response]<|im_end|>
<|im_start|>assistant
[Final response]<|im_end|>
Special Tokens
Token
ID
Purpose
<|im_start|>
32000
Start of message
<|im_end|>
32001
End of message
<thought>
32002
Start of reasoning
</thought>
32003
End of reasoning
<tool_call>
32004
Start of tool call
</tool_call>
32005
End of tool call
<tools>
32006
Start of tool definitions
</tools>
32007
End of tool definitions
Training Loss Curve
Step
Loss
PPL
10
6.72
825
100
2.15
8.6
400
1.06
2.9
730 (final)
0.90
2.5
Limitations
Optimized for English tool-calling; Korean support is limited
117M parameters - suitable for edge deployment but less capable than larger models
Best with structured tool-calling format; may struggle with free-form conversation