theabbie's picture
Publish Needle Pi coding-agent fine-tune
54c5de0 verified
|
Raw
History Blame Contribute Delete
4.65 kB
---
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](https://huggingface.co/Cactus-Compute/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`](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](https://github.com/cactus-compute/needle), install it, then download this repository:
```bash
hf download theabbie/needle-pi-coding-agent --local-dir needle-pi-coding-agent
```
```python
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`](integration/pi_tool_router.py) accepts this JSON on standard input:
```json
{"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](https://huggingface.co/Cactus-Compute/needle) and the [Needle source repository](https://github.com/cactus-compute/needle).