Needle：26M 参数模型也能用 Tool Calling

一个 26M 参数的"简单注意力网络"，用于函数调用，你甚至可以在 Mac/PC 上本地微调。在生产环境中，Needle 运行在 Cactus 上，预填充速度达到 6000 toks/sec，解码速度 1200。权重完全开源在 Cactus-Compute/needle，数据集生成也一样。

d=512, 8H/4KV, BPE=8192
                                  ┌──────────────┐
                                  │  Tool Call   │
                                  └──────┬───────┘
                                        ┌┴──────────┐
                                        │  Softmax  │
                                        └─────┬─────┘
                                        ┌─────┴─────┐
                                        │ Linear (T)│  ← tied
                                        └─────┬─────┘
                                        ┌─────┴─────┐
                                        │ ZCRMSNorm │
                                        └─────┬─────┘
                                     ┌────────┴────────┐
                                     │ Decoder x 8     │
                                     │┌───────────────┐│
                                     ││ ZCRMSNorm     ││
                                     ││ Masked Self   ││
                                     ││ Attn + RoPE   ││
                                     ││ Gated Residual││
                                     │├───────────────┤│
  ┌──────────────┐                   ││ ZCRMSNorm     ││
  │ Encoder x 12 │──────────────────────▶Cross Attn   ││
  │              │                   ││ Gated Residual││
  │ ┌──────────┐ │                   │└───────────────┘│
  │ │ZCRMSNorm │ │                   └────────┬────────┘
  │ │Self Attn │ │                      ┌─────┴─────┐
  │ │ GQA+RoPE │ │                      │ Embedding │  ← shared
  │ │Gated Res │ │                      └─────┬─────┘
  │ │          │ │                    ┌───────┴───────-┐
  │ │ (no FFN) │ │                    │[EOS]<tool_call>│
  │ └──────────┘ │                    │ + answer       │
  │              │                    └───────────────-┘
  └──────┬───────┘
         │
    ┌────┴──────┐
    │ Embedding │
    └────┬──────┘
         │
    ┌────┴──────┐
    │   Text    │
    │  query    │
    └───────────┘

在 16 个 TPU v6e 上预训练 200B tokens（27 小时）。

在 2B tokens 的单轮函数调用数据集上进行后训练（45 分钟）。

Needle 是简单注意力网络的一个实验性尝试，旨在为消费设备（手机、手表、眼镜...）重新定义小型 AI。虽然它在个人 AI 的单轮函数调用上超越了 FunctionGemma-270m、Qwen-0.6B、Granite-350m 和 LLaMA2.5-350m，但那些模型有更多功能/容量，在对话场景中表现更好。另外，小模型可能比较挑剔。请在下一部分使用 UI 在你自己的工具上测试，并点击按钮进行微调。

git clone https://github.com/cactus-compute/needle.git
cd needle && source ./setup
needle playground

在 http://127.0.0.1:7860 打开一个 web UI，你可以在其上测试和微调你自己的工具。权重会自动下载。

from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer

params, config = load_checkpoint("checkpoints/needle.pkl")
model = SimpleAttentionNetwork(config)
tokenizer = get_tokenizer()

result = generate(
    model, params, tokenizer,
    query="What's the weather in San Francisco?",
    tools='[{"name":"get_weather","description":"Get current weather for a city.","parameters":{"location":{"type":"string","description":"City name.","required":true}}}]',
    stream=False,
)
print(result)
# [{"name":"get_weather","arguments":{"location":"San Francisco"}}]

# Playground (generates data via Gemini, trains, evaluates, bundles result)
needle playground

# CLI (auto-downloads weights if not local)
needle finetune data.jsonl

JSONL 文件中的每一行都有三个字段：query、tools 和 answers。

{
  "name": "get_weather",
  "description": "Get current weather for a city.",
  "parameters": {
    "location": { "type": "string", "description": "City name.", "required": true }
  }
}

{ "name": "get_weather", "arguments": { "location": "Paris" } }

完整 JSONL 示例（每一行是一个训练样本，tools 和 answers 是 JSON 编码的字符串）：

{"query": "What's the weather in Paris?", "tools": "[{\"name\":\"get_weather\",\"description\":\"Get current weather for a city.\",\"parameters\":{\"location\":{\"type\":\"string\",\"description\":\"City name.\",\"required\":true}}}]", "answers": "[{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\"}}]"}
{"query": "Turn off the lights", "tools": "[{\"name\":\"get_weather\",\"description\":\"Get current weather for a city.\",\"parameters\":{\"location\":{\"type\":\"string\",\"description\":\"City name.\",\"required\":true}}},{\"name\":\"toggle_lights\",\"description\":\"Toggle smart lights on or off.\",\"parameters\":{\"state\":{\"type\":\"string\",\"description\":\"on or off.\",\"required\":true}}}]", "answers": "[{\"name\":\"toggle_lights\",\"arguments\":{\"state\":\"off\"}}]"}

每个工具至少提供 120 个示例（100 个训练 / 10 个验证 / 10 个测试）。示例太少会过拟合 —— 你会看到完美的训练指标，但模型不会泛化。改变 query 的表述方式，并包含多个工具可用的示例。

使用微调后的模型

微调会把最好的检查点保存为 checkpoints/needle_finetuned_<id>_best.pkl：

needle run --checkpoint checkpoints/needle_finetuned_*_best.pkl \
  --query "What's the weather?" --tools '[{"name":"get_weather","description":"Get current weather for a city.","parameters":{"location":{"type":"string","description":"City name.","required":true}}}]'

params, config = load_checkpoint("checkpoints/needle_finetuned_<id>_best.pkl")
model = SimpleAttentionNetwork(config)
result = generate(model, params, get_tokenizer(), query="...", tools='[...]', stream=False)

needle playground                  Test and finetune via web UI
needle finetune <data.jsonl>       Finetune on your own data
needle run --query "..." --tools   Single inference
needle train                       Full training run
needle pretrain                    Pretrain on PleIAs/SYNTH
needle eval --checkpoint <path>    Evaluate a checkpoint
needle tokenize                    Tokenize dataset
needle generate-data               Synthesize training data via Gemini
needle tpu <action>                TPU management (see docs/tpu.md)

@misc{ndubuaku2026needle,
  title={Needle},
  author={Henry Ndubuaku, Jakub Mroz,  Karen Mosoyan, Roman Shemet, Parkirat Sandhu, Satyajit Kumar, Noah Cylich, Justin H. Lee},
  year={2026},
  url={https://github.com/cactus-compute/needle}
}