Needle
<img src="https://raw.githubusercontent.com/cactus-compute/needle/HEAD/assets/banner.png" alt="Logo" style="border-radius: 30px; width: 100%;">
우리는 Gemini 3.1을 Mac/PC에서 로컬로 파인튜닝 (Finetune)할 수 있는 26M 파라미터 규모의 "Simple Attention Network"로 증류 (Distilled)했습니다.
실제 운영 환경에서 Needle은 Cactus 위에서 초당 6000 토큰 (toks/sec)의 프리필 (Prefill) 속도와 1200 토큰의 디코드 (Decode) 속도로 실행됩니다.
가중치 (Weights)와 데이터셋 생성 방법은 Cactus-Compute/needle에 완전히 공개되어 있습니다.
d=512, 8H/4KV, BPE=8192
┌──────────────┐
│ Tool Call │
...
- 16개의 TPU v6e에서 200B 토큰으로 사전 학습 (Pretrained) (27시간).
- 2B 토큰의 싱글샷 (Single-shot) 함수 호출 (Function call) 데이터셋으로 사후 학습 (Post-trained) (45분).
Needle는 소비자용 기기(휴대폰, 시계, 안경 등)를 위한 초소형 AI (Tiny AI)를 재정의하기 위해 설계된 Simple Attention Networks의 실험적 실행 결과입니다.
따라서 개인용 AI를 위한 싱글샷 함수 호출 성능에서는 FunctionGemma-270m, Qwen-0.6B, Graninte-350m, LFM2.5-350m을 능가하지만,
해당 모델들은 더 넓은 범위/용량 (Scope/Capacity)을 가지고 있으며 대화형 설정 (Conversational settings)에서 뛰어난 성능을 발휘합니다. 또한, 소형 모델은 까다로울 수 있습니다.
다음 섹션의 UI를 사용하여 직접 도구를 테스트하고, 버튼 클릭 한 번으로 그에 맞춰 파인튜닝 (Finetune)해 보시기 바랍니다.
빠른 시작 (Quickstart)
git clone https://github.com/cactus-compute/needle.git
cd needle && source ./setup
needle playground
http://127.0.0.1:7860 에서 웹 UI가 열리며, 여기서 자신만의 도구를 테스트하고 파인튜닝 (Finetune)할 수 있습니다. 가중치 (Weights)는 자동으로 다운로드됩니다.
사용법 (Usage - Python)
from needle import SimpleAttentionNetwork, load_checkpoint, generate, get_tokenizer
params, config = load_checkpoint("checkpoints/needle.pkl")
...
파인튜닝 (Finetuning)
JSONL 파일의 각 라인은 query, tools, answers 세 가지 필드를 가집니다.
도구 스키마 (Tool schema):
{
"name": "get_weather",
"description": "도시의 현재 날씨를 가져옵니다.",
...
응답 스키마 (Answer schema):
{ "name": "get_weather", "arguments": { "location": "Paris" } }
전체 JSONL 예시 (각 라인은 하나의 학습 예시이며, tools와 answers는 JSON 인코딩된 문자열입니다):
{"query": "파리 날씨가 어때?", "tools": "[{\"name\":\"get_weather\",\"description\":\"도시의 현재 날씨를 가져옵니다.\",\"parameters\":{\"location\":{\"type\":\"string\",\"description\":\"도시 이름.\",\"required\":true}}}]", "answers": "[{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\"}}]"}
{"query": "불을 꺼줘", "tools": "[{\"name\":\"get_weather\",\"description\":\"도시의 현재 날씨를 가져옵니다.\",\"parameters\":{\"location\":{\"type\":\"string\",\"description\":\"도시 이름.\",\"required\":true}}},{\"name\":\"toggle_lights\",\"description\":\"스마트 조명을 켜거나 끄는 기능을 전환합니다.\",\"parameters\":{\"state\":{\"type\":\"string\",\"description\":\"on 또는 off.\",\"required\":true}}}]", "answers": "[{\"name\":\"toggle_lights\",\"arguments\":{\"state\":\"off\"}}]"}
각 도구당 최소 120개의 예시를 제공해야 합니다 (학습 100개 / 검증 10개 / 테스트 10개). 예시가 적으면 과적합(overfit)될 수 있습니다. 완벽한 학습 지표가 나타나더라도 모델이 일반화되지 못할 수 있으니, 질의어 표현을 다양하게 하고 여러 도구가 사용 가능한 예시를 포함하세요.
파인튜닝된 모델 사용하기
파인튜닝은 최적의 체크포인트(checkpoint)를 checkpoints/needle_finetuned_<id>_best.pkl로 저장합니다:
needle run --checkpoint checkpoints/needle_finetuned_*_best.pkl \
--query "날씨가 어때?" --tools '[{"name":"get_weather","description":"도시의 현재 날씨를 가져옵니다.","parameters":{"location":{"type":"string","description":"도시 이름.","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)
CLI
needle playground 웹 UI를 통해 테스트 및 파인튜닝
needle finetune <data.jsonl> 자체 데이터를 이용한 파인튜닝
needle run --query "..." --tools 단일 추론
...
@misc{ndubuaku2026needle,
title={Needle},
author={Henry Ndubuaku, Jakub Mroz, Karen Mosoyan, Roman Shemet, Parkirat Sandhu, Satyajit Kumar, Noah Cylich, Justin H. Lee},
...