theabbie's picture
Publish Needle Pi coding-agent fine-tune
54c5de0 verified
|
Raw
History Blame Contribute Delete
4.65 kB
metadata
license: mit
library_name: jax
pipeline_tag: text-generation
base_model: Cactus-Compute/needle
language:
  - en
tags:
  - needle
  - function-calling
  - tool-use
  - coding-agent
  - pi-coding-agent
  - jax
  - flax
  - encoder-decoder
  - on-device
model-index:
  - name: Needle Pi Coding Agent
    results:
      - task:
          type: text-generation
          name: Pi tool calling
        dataset:
          name: Held-out Pi coding tools
          type: pi-tools-held-out
        metrics:
          - type: exact_match
            value: 0.8
            name: Exact tool-call match
          - type: f1
            value: 0.8
            name: Tool-call F1

Needle Pi Coding Agent

A 26.3M-parameter Cactus Needle fine-tune for routing coding-agent intents into Pi's four built-in tools: read, bash, edit, and write.

It is designed as a small parser/router, not a knowledge engine. An upstream planner should provide concrete paths, commands, content, or replacements; this model selects the tool and copies those arguments into a structured call.

Results

Evaluation used a held-out, balanced 40-example test set with 10 examples per tool.

Metric Base Needle This fine-tune
Exact match 27.5% 80.0%
Tool-call F1 28.95% 80.0%
Tool-name F1 88.0% 100.0%
JSON parse rate 82.5% 100.0%
Argument accuracy 32.35% 80.0%

Exact matches by tool improved from read 8/10, bash 2/10, edit 0/10, write 1/10 to read 9/10, bash 8/10, edit 9/10, write 6/10.

These are local single-run measurements, not an external benchmark. See metrics.json for machine-readable results.

Training

  • Base: Cactus-Compute/needle
  • Examples: 1,500 (bash: 450; read, edit, write: 350 each)
  • Split: 1,420 train / 40 validation / 40 test
  • Epochs: 1
  • Batch size: 32
  • Training steps: 22
  • Optimizer learning rates: Adam 3e-5, Muon 0.02
  • Checkpoint-selection validation exact match: 70.0%
  • Synthetic-data generator: Codex CLI, with strict schema, grounding, and deduplication validation

The additional bash data emphasizes parsing fully specified arbitrary shell commands, including quoting, flags, environment variables, pipelines, redirects, subshells, and multiline commands. The model is not expected to invent specialist commands from vague goals.

Usage

Clone the Needle repository, install it, then download this repository:

hf download theabbie/needle-pi-coding-agent --local-dir needle-pi-coding-agent
import json
from needle import SimpleAttentionNetwork, generate, get_tokenizer, load_checkpoint

params, config = load_checkpoint("needle-pi-coding-agent/needle-pi-coding-agent.pkl")
model = SimpleAttentionNetwork(config)

tools = json.loads(open("needle-pi-coding-agent/training/pi_tools.json").read())
result = generate(
    model,
    params,
    get_tokenizer(),
    query="Run bash command: curl -s https://api.ipify.org",
    tools=json.dumps(tools, separators=(",", ":")),
    stream=False,
)
print(result)
# [{"name":"bash","arguments":{"command":"curl -s https://api.ipify.org"}}]

The bundled integration/pi_tool_router.py accepts this JSON on standard input:

{"intent":"Run bash command: curl -s https://api.ipify.org","tools":[...]}

and emits a JSON array containing the selected call.

Included artifacts

  • needle-pi-coding-agent.pkl: best native Needle/JAX checkpoint
  • config.json and tokenizer/: architecture and tokenizer files
  • training/pi_tools_1500.jsonl: full 1,500-example training corpus
  • training/pi_tools.json: flattened schemas for Pi's built-in tools
  • training/generate_pi_dataset.py: resumable, atomically persisted generator
  • training/codex_client.py: clean Codex CLI text-generation adapter
  • training/pi_tools_1500.progress.json: final generation manifest
  • integration/pi_tool_router.py: minimal stdin/stdout router used by the Pi provider prototype

Limitations

  • This is a focused single-shot tool router, not a conversational language model.
  • Exact edit extraction is strongest when the planner explicitly supplies the edits array with oldText and newText values.
  • Long or whitespace-sensitive write/edit payloads can still be copied imperfectly.
  • Shell commands are executed by the surrounding agent runtime; this model does not provide sandboxing or command safety.

License and attribution

MIT. Based on Cactus-Compute/needle and the Needle source repository.