Views
No views yet
ibm-granite/granite-4.1-3b that learns to pick the right function(s) to call from a list of available tools, given a natural-language query.selected_tools list.1{
2 "query": "What's the weather in Paris and the time in Tokyo?",
3 "tools": [
4 {"name": "get_weather", "description": "Get the current weather for a city."},
5 {"name": "get_time", "description": "Get the current local time for a city."},
6 {"name": "send_email", "description": "Send an email to a recipient."}
7 ]
8}{"selected_tools": ["get_weather", "get_time"]}chat_template):<|start_of_role|>system<|end_of_role|>You are a tool-selection assistant. Given a user query and a list of available tools, return the names of the tools that should be called.<|end_of_text|>
<|start_of_role|>user<|end_of_role|>{...JSON above...}<|end_of_text|>
<|start_of_role|>assistant<|end_of_role|>1from peft import PeftModel
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4base = AutoModelForCausalLM.from_pretrained("ibm-granite/granite-4.1-3b", torch_dtype="bfloat16", device_map="auto")
5tokenizer = AutoTokenizer.from_pretrained("barha/granite-4.1-3b-tool-selector")
6model = PeftModel.from_pretrained(base, "barha/granite-4.1-3b-tool-selector")
7model.eval()
8
9SYSTEM = "You are a tool-selection assistant. Given a user query and a list of available tools, return the names of the tools that should be called."
10user_payload = {
11 "query": "What's the weather in Paris?",
12 "tools": [{"name": "get_weather", "description": "Get the current weather for a city."}],
13}
14import json
15prompt = (
16 f"<|start_of_role|>system<|end_of_role|>{SYSTEM}<|end_of_text|>\n"
17 f"<|start_of_role|>user<|end_of_role|>{json.dumps(user_payload)}<|end_of_text|>\n"
18 f"<|start_of_role|>assistant<|end_of_role|>"
19)
20inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
21out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
22gen = tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False)
23print(gen.split("<|end_of_text|>", 1)[0].strip())Salesforce/xlam-function-calling-60k (gated dataset on Hugging Face), filtered to keep only rows whose answers list is non-empty (no-call abstentions are excluded). 90/10 train/val split, deterministic (no shuffle), seed 42.tools list is reduced to {name, description} pairs (parameters dropped — the adapter routes by name only) and concatenated with the query in the user turn.ibm-granite/granite-4.1-3b (bf16, ungated, Apache 2.0)r=8, alpha=16, dropout=0.05, bias="none"q_proj, k_proj, v_proj, o_projlr=2e-4, cosine schedule, warmup ratio 0.03| epoch | step | train loss | eval loss |
|---|---|---|---|
| 0.41 | 1400 | 0.55 | 0.554 |
| 0.95 | 3200 | 0.39 | — |
| 1.41 | 4760 | 0.38 | — |
| 2.03 | 6850 | 0.32 | — |
| 2.84 | 9580 | 0.33 | 0.383 |
| 2.96 | 10000 | 0.31 | 0.383 |
xlam-function-calling-60k used during training; the adapter never saw these examples). Predictions parsed from the JSON selected_tools field, scored as sets of tool names.| metric | value |
|---|---|
| exact set match | 0.9930 (5,958 / 6,000) |
| macro F1 | 0.9960 |
| precision | 0.9963 |
| recall | 0.9960 |
| parse failure rate | 0.0007 (4 / 6,000) |
max_new_tokens=128, greedy).train/tool_selector/eval_adapter.py and the train/jobs/tool-selector-eval.yaml AppWrapper.