Needle 2 概述

Needle 2 是一个面向工具调用、设备控制与结构化抽取的开源 45M 参数模型。整个模型被打包为约 14MB 的二进制单文件，运行一个完整会话仅占用约 28MB 内存。它基于团队提出的 Simple Attention Network 架构，使用 Cactus Quants 量化方案压缩到 CQ2-bit，并与自有推理引擎一同发布。在多项基准测试中，Needle 2 以 5 到 70 倍的体积优势与 FunctionGemma 270M、LFM2.5 230M、Apple FM 等小型模型表现相当，而量化精度为 2 bit，对比通常的 f16 显著更小。

项目定位与核心特性

• 自包含：权重内嵌于单个引擎，无额外模型文件，推理过程无网络请求。
• 简洁接口：文本输入、JSON 输出；根据工具 schema 编译字节级语法（grammar），每个 token 都受其约束。
• 置信度门控：每个响应附带由独立学习头输出的校准置信度分数，可设置阈值决定执行还是上报。
• 工具检索：声明一个大型工具目录，内置检索头每轮只渲染前五个最相关工具，并将语法限制在该子集内。
• 有界内存：使用 256 token 滑动窗口，并将工具固定为 KV sink，使总占用在长对话中保持稳定在 28MB 附近。

技术架构：Simple Attention Network

每个网络块内部更新规则如下：

• 使用 RMS 归一化的四路残差流拼接作为输入 x̂。
• 通过固定的 Walsh-Hadamard 正交变换 H 替换传统 FFN，矩阵无需参数、可在 O(n log n) 时间内应用。
• 通过 n-gram 哈希表取出 engram 键值对 (kₜ, vₜ)。
• 路由 logits A 经 Sinkhorn 迭代得到双随机归一化矩阵 P。
• 所有门控参数（a、b、g 以及各 σ-gate）均为可学习且输入相关。
• 注意力与 MLP 残差均经过 sandwich-norm 并由门控调节；engram 仅在两层触发；解码阶段通过字节级语法约束输出。

完整的架构论文详见 arXiv:2607.18363。

快速上手

安装命令：

pip install cactus-needle

工具描述质量直接决定模型调用效果。下面是核心用法示例。

简单工具调用：使用装饰器定义函数，函数签名为参数提供类型信息，docstring 作为工具描述；run() 会自动完成调用循环并返回结果。

import needle

@needle.tool
def get_weather(city: str):
    """获取指定城市的当前天气。"""
    return {"city": city, "temp_c": 27, "sky": "clear"}

agent = needle.Needle(tools=[get_weather])
print(agent.run("Lagos 现在天气如何？")["results"])
# [{'city': 'Lagos', 'temp_c': 27, 'sky': 'clear'}]

结构化抽取：传入 Pydantic 模型即可获得类型化对象。

from pydantic import BaseModel

class Invoice(BaseModel):
    vendor: str
    total: float
    due_date: str

invoice = needle.extract("Invoice from Acme Corp, $1,200.00, due 2026-09-01", Invoice)

Playground

可在浏览器中直接试用任意模型：选择预设、编辑工具或提示词后运行，后续追问延续同一会话。

needle playground                      # 默认模型，监听 127.0.0.1:7860
needle playground --weights my.cact    # 使用自定义权重

UI 中的 Finetune on these tools 按钮会直接触发下文的微调流程，并返回可下载的 .cact 文件。

微调流程

Needle 在冻结基座上做 LoRA 微调，导出时合并适配器，最终仍是单文件 .cact，可在同一引擎直接运行。

数据格式：JSONL，每行一个样本；reasoning 可选，与工具无关的样本使用 answers: []。

{"query": "dim the kitchen to 10", "tools": [{"name": "set_lights", "parameters": {"type": "object", "properties": {"room": {"type": "string"}, "brightness": {"type": "integer"}}, "required": ["room"]}}], "answers": [{"name": "set_lights", "arguments": {"room": "kitchen", "brightness": 10}}], "reasoning": "'kitchen' -> room; 'dim to 10' -> brightness 10"}

1. 合成数据（可选）：需设置 OPENROUTER_API_KEY。可从工具 schema 文件生成种子样本，也可扩展现有数据集。

export OPENROUTER_API_KEY=sk-or-...
needle generate-data --tools my_tools.json --num-samples 500 --output data.jsonl
needle generate-data --augment data.jsonl --num-samples 500      # 扩展已有 JSONL

通过设置 OPENROUTER_URL 可切换为 OpenAI 兼容网关。

2. LoRA 微调：基座检查点如未通过 --checkpoint 指定，会自动从 Hugging Face 下载。--generate N 会先用工具 schema 再合成 N 条样本（同样需要 API key）。

needle finetune data.jsonl --epochs 10
needle finetune data.jsonl --epochs 10 --generate 300 --lora-rank 16 --lora-alpha 32

常用参数：--epochs（默认 3）、--lora-rank（16）、--lora-alpha（32）、--lr（1e-4）、--batch-size（16）、--max-len（1024）、--val-split（0.1）、--checkpoint <base.pkl>、--out <adapter.pkl>。适配器默认输出到 checkpoints/needle_lora.pkl，每轮打印验证损失。

训练基于 JAX，可在任意支持的后端运行。NVIDIA 平台启用 CUDA：

pip install "cactus-needle[gpu]"

Apple Silicon 启用 Metal：

pip install "cactus-needle[metal]"

3. 打包为 .cact：合并适配器并量化。基座缺失会自动下载。

needle build checkpoints/needle2.pkl --lora checkpoints/needle_lora.pkl --out my_needle.cact

--bits 2 可进一步压缩；默认按检查点声明的逐层位宽导出，若未声明则回退为 4。设置 NEEDLE_HF_REPO=<user>/<model> 并加上 --upload 可发布到 Hugging Face。对应的 needle download <user>/<model>/my_needle.cact 用于拉取已发布的 .cact；needle download macos-arm64 等命令可获取对应平台的引擎运行器。

4. 运行微调后模型：引擎与权重解耦，因此自定义 .cact 可直接被加载，无需重新编译。

import needle
agent = needle.Needle(weights="my_needle.cact", tools=[...])
agent.run("...")

引用

Needle 2 由 Cactus Compute 团队构建，使用 BibTeX 引用如下：

@misc{needle2_2026,
  title        = {Needle 2: A 45M-Parameter Foundation Tool-Calling Model for Tiny Devices},
  author       = {Ndubuaku, Henry and Mosoyan, Karen and Mroz, Jakub and Cylich, Noah and
                  Kumar, Satyajit and Sandhu, Parkirat and Shemet, Roman and Lee, Justin H.},
  year         = {2026},
  organization = {Cactus Compute, Inc.},
  howpublished = {\url{https://github.com/cactus-compute/needle}}
}