Needle 2

Needle 2 是一个开放的 45M 参数模型，专为工具调用、设备操控和结构化信息提取而设计。整个模型是一个仅 14MB 的单一二进制文件，运行完整会话仅需约 28MB 内存。它基于我们提出的 Simple Attention Network 研究成果构建，使用 Cactus Quants 压缩至 CQ2 位精度，并集成到其专属引擎中。在下方基准测试中，Needle 2 与 FunctionGemma 270M、LFM2.5 230M 和 Apple FM 等其他小型模型互有胜负，但体积小 5 到 70 倍，且仅用 2 位精度对比它们的 f16 精度。

本仓库为 Python 包，涵盖推理、LoRA 微调和导出功能。执行 pip install cactus-needle，描述你的工具，即可在 Python 中调用。推理引擎首次从 Hugging Face 获取并缓存，无需额外构建；离线部署（适用于气隙设备）的说明见 doc/apis.md。

• 自包含：权重内置于单个 14MB 引擎中，无需管理单独的模型文件，推理过程不依赖网络。
• 简洁接口：工具调用以结构化数据返回——文本输入，JSON 输出；由你的模式（schema）编译而成的字节级语法约束每一个 token。
• 置信度门控：每次响应都附带来自学习头的校准置信度分数；设定阈值，高于阈值则执行，低于阈值则升级处理。
• 工具检索：声明一个大型工具目录，内置检索头每次仅呈现排名前五的工具，语法亦被约束在该子集内。
• 有界内存：256-token 滑动窗口，工具作为 KV 锚点固定其中，无论对话多长，总内存始终保持在约 28MB。

权重：huggingface.co/Cactus-Compute/needle2 · 源码：github.com/cactus-compute/needle。

Simple Attention Network

Needle 2 是一个 Simple Attention Network，这是我们针对密集型小模型的设计方案：以 Hadamard MLP 替代 FFN，采用 GQA 注意力、engram 键值记忆和多通道超连接。详细设计与消融实验见论文：arXiv:2607.18363。

每个模块都携带其更新规则。其中 x̂ 是四条残差流的 RMS 归一化展平结果，H 是正交 Walsh-Hadamard 变换（固定矩阵，以 n log n 时间复杂度计算，无需读取权重），(kₜ, vₜ) 行从哈希 n-gram 表中收集，P 是通过 Sinkhorn 迭代计算的路由 logits A 的双随机归一化；a、b、g 及所有 σ 门均为可学习且依赖输入的参数。注意力与 MLP 残差均采用 sandwich 归一化和门控，engram 位点在两层触发，解码过程受从声明模式编译的字节级语法约束。

快速开始

pip install cactus-needle

Needle 会读取你的工具描述来决定调用什么以及如何填写参数，因此把描述写好才是关键。

简单：装饰一个函数即可。函数签名提供参数类型，docstring 即工具描述，run() 完成整个闭环：模型选择调用方式，Needle 执行你的函数，将结果回传，并返回最终响应，同时将已执行的工具结果作为 results 附上。

import needle

@needle.tool
def get_weather(city: str):
    "Get the current weather for a city."
    return {"city": city, "temp_c": 27, "sky": "clear"}

agent = needle.Needle(tools=[get_weather])
print(agent.run("what's it like in Lagos right now?")["results"])
# [{'city': 'Lagos', 'temp_c': 27, 'sky': 'clear'}]

信息提取：要从文本中提取结构化数据，只需声明数据结构并调用 extract() 方法。传入一个 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)
print(invoice.vendor, invoice.total)   # -> Acme Corp 1200.0

参数描述与选项、编译进解码语法的值约束、原始 JSON 模式、通过 complete() 驱动循环、响应契约、系统事实、工具检索以及置信度门控，均已在 doc/apis.md 中涵盖。

Playground

在浏览器中试用任意模型：选择一个预设，编辑工具或提示词，然后点击 Run。后续查询将继续同一段对话。

needle playground                      # base model, http://127.0.0.1:7860
needle playground --weights my.cact    # a tuned model

服务器在提供服务之前会先下载并初始化模型，因此首次查询即刻响应。在这些工具上进行微调按钮从界面运行下述微调流程，并返回一个可下载的 .cact 文件。

微调

Needle 在冻结的基础模型上使用 LoRA 进行微调，并在导出时合并适配器，因此运行成本低，且微调后的模型仍然是单个可运行在同一引擎上的 .cact 文件。流程为：（可选）合成数据、LoRA 微调，然后构建微调后的 .cact。有关数据集规模、损失曲线解读及故障排查，请参阅 doc/finetuning.md。

数据格式。 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。从工具模式文件（tool schema file）中生成种子数据，或扩展现有数据集：

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      # expand an existing JSONL

将 OPENROUTER_URL 设置为使用兼容 OpenAI 的网关，以替代默认的 OpenRouter 端点。

2. LoRA 微调。 如果你不传入 --checkpoint 参数，基础检查点会自动从 Hugging Face 下载。--generate N 会先从你数据中的工具里合成 N 个额外示例（同样需要 OPENROUTER_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。每个 epoch 结束后会从保留的验证集中输出验证损失。

训练过程基于纯 JAX 实现，可在 JAX 支持的任何加速器上运行。在 NVIDIA 机器上安装 CUDA 版本后，同样的命令即可在 GPU 上完成训练：

pip install "cactus-needle[gpu]"

在 Apple Silicon 上，metal 附加组件可在 GPU 上进行训练：

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=<you>/<model> 并传入 --upload 以发布 .cact 文件。对应的 needle download <you>/<model>/my_needle.cact 命令可在任意机器上拉取已发布的归档文件。

4. 运行它。 该引擎与权重无关，因此调优后的 .cact 文件可直接在其上运行——无需重新编译：

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

引用

Needle 2 由 Cactus Compute 团队开发。若您在研究中使用了本项目，请引用：

@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}}
}

如需洽谈合作、协同共赢，或希望将Needle2部署到您的产品中，欢迎通过founders@cactuscompute.com与我们联系。