cactus-needle-toolcall-lora

A LoRA fine-tune of Cactus-Compute/needle2 — Needle 2, the 45M-parameter tool-calling model — for an 8-tool home-automation suite. It picks the right tool more often than the stock base model and keeps the base's refusal behaviour intact, all while staying a single 23MB .cact that runs on the unmodified engine. No new runtime, no recompilation.

Source + research: the full codebase — dataset generator, Kaggle TPU training pipeline, and benchmark harness — lives at github.com/instax-dutta/cactus-needle-toolcall-lora. An in-depth field report on the whole project — data engineering, the schema-mismatch root cause behind a mode: "bake" artifact, refusal balancing, engine quirks, and benchmark methodology — is published as a GitHub Gist.

Who this is for: developers shipping a Needle-based assistant against a fixed set of domain tools — a smart-home hub, an on-device agent, an edge appliance — who want better tool selection than the stock base model without graduating to a larger model or building a custom engine. Runs wherever the base model runs: Raspberry Pi, phones, Apple Silicon, WebAssembly, and the rest of the supported platforms.

Why LoRA fine-tuning instead of prompt engineering

Needle's engine already guarantees well-formed calls — a byte-level grammar compiled from your schemas constrains every token, so the JSON cannot be malformed. What the grammar cannot decide is which tool to call and which values to fill, and that is exactly where a 45M base shows its size: an under-described phrasing gets routed to the wrong tool, or an argument value gets guessed instead of omitted.

Fine-tuning fixes that directly. LoRA, rather than a full fine-tune, trains a small set of adapter weights on top of the untouched base — so the model keeps everything Needle 2 learned at pretraining and only re-weights the tool-calling behaviour you care about. The adapter merges through the same Cactus Quant pipeline (needle build) and ships as a .cact the stock engine loads as-is.

Training

Parameter	Value
Base model	Cactus-Compute/needle2 (45M)
Method	LoRA (needle finetune), merged and exported to Cactus Quants
LoRA rank / alpha	32 / 32
Epochs	4
Max context (--max-len)	768
Batch size	8
Learning rate	1e-4
Precision	float32
Hardware	Google TPU v5e-8 (Kaggle)
Final loss	0.0274

Dataset

462 hand-generated examples (seed 42) over the 8-tool home-automation suite — tool phrasings, argument phrasings, and refusals in roughly a 70/30 split:

Tool	Examples
set_thermostat	55
play_music	55
set_alarm	40
send_message	40
set_lights	35
get_weather	35
add_to_calendar	35
lock_doors	30
refusals (empty call)	137

our_data.jsonl is included in this repo — same format the base repo's finetune consumes, so it doubles as a starter set for teaching Needle your own tool vocabulary.

Performance

Held-out benchmark (16 tool calls + 2 off-topic refusals), measured end-to-end through the shipped engine, strict name + argument match, stable across 3 runs:

Metric	This model	Needle 2 (base)
Tool name selected correctly	16/16	14/16
Full call (name + every argument)	13/16	12/16
Off-topic input refused (empty call)	2/2	2/2

The finetune recovers the two tools the base mis-routes — the base sends "message Alex that dinner is ready" to get_weather, and goes silent on a calendar request — and is perfect on tool selection. The remaining argument misses are shared formatting quirks rather than new regressions: "put on lo-fi beats" extracts query: "lo-fi beats" instead of "lo-fi" (the base does the same), and the base's 6:30am is normalized to 06:30.

Honest caveats: 18 held-out cases is a small benchmark, and the base model is already strong on this task — the real win is tool selection (16/16 vs 14/16), not raw capability. Treat the full-match delta of 13 vs 12 as directional, and re-benchmark against your own schemas before trusting it in production.

Files

File	What it is	Size
our_needle.cact	Ship-ready model (4-bit W4A8 Cactus Quant)	23 MB
checkpoints/needle_lora.pkl	LoRA adapter weights	16 MB
our_data.jsonl	Training data (462 examples)	0.9 MB

Quickstart

from huggingface_hub import hf_hub_download
import needle, json

weights = hf_hub_download("saidutta69/cactus-needle-toolcall-lora", "our_needle.cact")

tools = json.load(open("tools.json"))   # your home-automation schemas
agent = needle.Needle(weights=weights, tools=tools)

agent.complete("make it 19 degrees heat")
# {'type': 'call', 'function_calls': [{'name': 'set_thermostat', 'arguments': {'temperature': 19, 'mode': 'heat'}}], ...}

The engine is weights-agnostic — load any tuned .cact with the stock cactus-needle package (2.0.1) and its fetched engine, no recompilation. pip install cactus-needle, describe your tools, and call them.

Responsible use

This is a tool router, not a chatbot. It answers exclusively with function calls against the tools you declare; anything no declared tool can serve is refused with the empty call. It inherits Needle 2's behaviour contract — argument values are evidence-based, and refusal quality depends on your tool descriptions being complete. Benchmark on your own schema before production.

License

Apache-2.0, inherited from the base Cactus-Compute/needle2.

Made with ❤️ by RACER IS OP