cactus-compute/needle

# Needle 我们将 Gemini 3.1 蒸馏成了一个拥有 2600 万参数的“[Simple Attention Network](docs/simple_attention_networks.md)”，你甚至可以在你的 Mac/PC 上对其进行本地微调。 在生产环境中，Needle 运行在 [Cactus](https://github.com/cactus-compute/cactus) 上，预填充速度为 6000 toks/sec，解码速度为 1200。 模型权重及数据集生成过程已在 [Cactus-Compute/needle](https://huggingface.co/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]│ │ └──────────┘ │ │ + answer │ │ │ └───────────────-┘ └──────┬───────┘ │ ┌────┴──────┐ │ Embedding │ └────┬──────┘ │ ┌────┴──────┐ │ Text │ │ query │ └───────────┘ ``` - 在 16 个 TPU v6e 上预训练了 200B tokens（耗时 27 小时）。 - 在 2B tokens 的单次函数调用数据集上进行了后训练（耗时 45 分钟）。 Needle 是 Simple Attention Networks 的一次实验性尝试，旨在重新定义适用于消费级设备（手机、手表、眼镜……）的微型 AI。 因此，尽管在面向个人 AI 的单次函数调用方面，它击败了 FunctionGemma-270m、Qwen-0.6B、Graninte-350m 和 LFM2.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，你可以在其中针对自己的工具进行测试和微调。权重会自动下载。 ## 使用方法（Python） ``` 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（通过 Gemini 生成数据、进行训练、评估并打包结果） needle playground # CLI（若本地不存在则自动下载 weights） needle finetune data.jsonl ``` ### 数据格式 JSONL 文件中的每一行包含三个字段：`query`、`tools` 和 `answers`。 **Tool schema：** ``` { "name": "get_weather", "description": "Get current weather for a city.", "parameters": { "location": { "type": "string", "description": "City name.", "required": true } } } ``` **Answer schema：** ``` { "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\"}}]"} ``` 为每个 tool 至少提供 **120 个示例**（100 个训练 / 10 个验证 / 10 个测试）。示例太少会导致过拟合——你会看到完美的训练指标，但模型无法泛化。请丰富查询的措辞，并包含具有多个可用 tool 的示例。 ### 使用微调后的模型 微调过程会将最佳检查点保存为 `checkpoints/needle_finetuned__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__best.pkl") model = SimpleAttentionNetwork(config) result = generate(model, params, get_tokenizer(), query="...", tools='[...]', stream=False) ``` ## CLI ``` needle playground Test and finetune via web UI needle finetune 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 Evaluate a checkpoint needle tokenize Tokenize dataset needle generate-data Synthesize training data via Gemini needle tpu 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} } ```

标签：DLL 劫持, IaC 扫描, 人工智能, 函数调用, 大语言模型, 模型压缩, 用户模式Hook绕过, 轻量级模型, 边缘计算, 逆向工具