Inputs and outputs

Input (user turn, JSON-encoded):

json
Copy code
{
  "query": "What's the weather in Paris and the time in Tokyo?",
  "tools": [
    {"name": "get_weather", "description": "Get the current weather for a city."},
    {"name": "get_time",    "description": "Get the current local time for a city."},
    {"name": "send_email",  "description": "Send an email to a recipient."}
  ]
}

Output (assistant turn, JSON):

json
Copy code
{"selected_tools": ["get_weather", "get_time"]}

The full prompt uses Granite's role-tagged tokens directly (the base tokenizer has no chat_template):

markdown
Copy code
<|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|>

How to use

python
Copy code
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained("ibm-granite/granite-4.1-3b", torch_dtype="bfloat16", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("barha/granite-4.1-3b-tool-selector")
model = PeftModel.from_pretrained(base, "barha/granite-4.1-3b-tool-selector")
model.eval()

SYSTEM = "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."
user_payload = {
    "query": "What's the weather in Paris?",
    "tools": [{"name": "get_weather", "description": "Get the current weather for a city."}],
}
import json
prompt = (
    f"<|start_of_role|>system<|end_of_role|>{SYSTEM}<|end_of_text|>\n"
    f"<|start_of_role|>user<|end_of_role|>{json.dumps(user_payload)}<|end_of_text|>\n"
    f"<|start_of_role|>assistant<|end_of_role|>"
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
gen = tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False)
print(gen.split("<|end_of_text|>", 1)[0].strip())

Training data

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.

For each row, the JSON-encoded 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.

After filtering: 54,000 train / 6,000 validation examples.

Training procedure

• Base model: ibm-granite/granite-4.1-3b (bf16, ungated, Apache 2.0)
• Adapter: LoRA, r=8, alpha=16, dropout=0.05, bias="none"
• Target modules: q_proj, k_proj, v_proj, o_proj
• Trainable parameters: ~5.2 M (~0.17 % of base)
• Epochs: 3
• Optimizer: AdamW, lr=2e-4, cosine schedule, warmup ratio 0.03
• Effective batch size: 16 (4 GPUs × per-device 4 × accum 1)
• Max sequence length: 2048
• Mixed precision: bf16
• Gradient checkpointing: non-reentrant
• Hardware: 4× NVIDIA A100 80GB (single node)
• Wall-clock: ~78 min training (post-training eval was abandoned)

Loss curve

Eval loss dropped from 0.554 → 0.383 with no signs of overfitting at the end of epoch 3.

Evaluation

Greedy generation on the held-out 6,000-example val split (same 90/10 deterministic slice of 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.

Generation wall-time: ~15.4 min on 1× A100 (bs=8, max_new_tokens=128, greedy).

Reproduce with train/tool_selector/eval_adapter.py and the train/jobs/tool-selector-eval.yaml AppWrapper.

Caveat — in-distribution only. Train and val are both deterministic slices of the same dataset, so this is an upper bound under matched distribution. The numbers do not speak to OOD generalization (different tool taxonomies, ambiguous queries, adversarial tool lists). Run your own behavioral eval against your real tool set before relying on this adapter in production.