Cactus a open-sourcé Needle, un modèle d'appel de fonction de 26M paramètres conçu pour fonctionner sur des téléphones, montres et lunettes économiques. Il atteint 6000 tok/s en pré-remplissage et 1200 tok/s en décodage sur des appareils grand public grâce à son moteur d'inférence personnalisé, Cactus.

Architecture : Simple Attention Networks

Needle utilise un Simple Attention Network — pas de MLP nulle part. Le modèle entier se compose de couches d'attention et de gating. Conception clé : d=512, 8H/4KV, BPE=8192, avec une structure encodeur-décodeur (12 couches encodeur, 8 couches décodeur) utilisant l'attention croisée, l'auto-attention masquée avec RoPE, et des embeddings liés.

Détails d'entraînement

• Pré-entraîné sur 200B tokens avec 16 TPU v6e (27 heures)
• Post-entraîné sur 2B tokens de données synthétisées d'appel de fonction (45 minutes)
• Données synthétisées via Gemini avec 15 catégories d'outils (minuteries, messagerie, navigation, maison intelligente, etc.)

Résultats des benchmarks

Needle bat FunctionGemma-270M, Qwen-0.6B, Granite-350M et LFM2.5-350M sur l'appel de fonction en un seul tour. Cependant, ces modèles ont plus de portée/capacité et excellent dans les contextes conversationnels.

Démarrage rapide

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

Ouvre une interface web à http://127.0.0.1:7860 pour tester et affiner vos propres outils.

Utilisation (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="Quel temps fait-il à San Francisco ?",
    tools='[{"name":"get_weather","parameters":{"location":"string"}}]',
    stream=False
)
print(result)
[{"name":"get_weather","arguments":{"location":"San Francisco"}}]

Affinage local

# via playground (génère automatiquement des données via Gemini)
needle playground
ou fournissez vos propres données
needle finetune data.jsonl

Disponibilité

Les poids sont sur Hugging Face : Cactus-Compute/needle. Tout est sous licence MIT.