Publish Needle Pi coding-agent fine-tune
Browse filesNative Needle checkpoint, tokenizer, model card, 1,500-example corpus, reproducible generator, metrics, and Pi router integration.
- LICENSE +21 -0
- README.md +128 -0
- config.json +5 -0
- integration/pi_tool_router.py +35 -0
- metrics.json +23 -0
- needle-pi-coding-agent.pkl +3 -0
- tokenizer/needle.model +3 -0
- tokenizer/needle.vocab +0 -0
- training/codex_client.py +58 -0
- training/generate_pi_dataset.py +291 -0
- training/pi_tools.json +35 -0
- training/pi_tools_1500.jsonl +0 -0
- training/pi_tools_1500.progress.json +23 -0
LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 Cactus Compute
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
README.md
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: mit
|
| 3 |
+
library_name: jax
|
| 4 |
+
pipeline_tag: text-generation
|
| 5 |
+
base_model: Cactus-Compute/needle
|
| 6 |
+
language:
|
| 7 |
+
- en
|
| 8 |
+
tags:
|
| 9 |
+
- needle
|
| 10 |
+
- function-calling
|
| 11 |
+
- tool-use
|
| 12 |
+
- coding-agent
|
| 13 |
+
- pi-coding-agent
|
| 14 |
+
- jax
|
| 15 |
+
- flax
|
| 16 |
+
- encoder-decoder
|
| 17 |
+
- on-device
|
| 18 |
+
model-index:
|
| 19 |
+
- name: Needle Pi Coding Agent
|
| 20 |
+
results:
|
| 21 |
+
- task:
|
| 22 |
+
type: text-generation
|
| 23 |
+
name: Pi tool calling
|
| 24 |
+
dataset:
|
| 25 |
+
name: Held-out Pi coding tools
|
| 26 |
+
type: pi-tools-held-out
|
| 27 |
+
metrics:
|
| 28 |
+
- type: exact_match
|
| 29 |
+
value: 0.8
|
| 30 |
+
name: Exact tool-call match
|
| 31 |
+
- type: f1
|
| 32 |
+
value: 0.8
|
| 33 |
+
name: Tool-call F1
|
| 34 |
+
---
|
| 35 |
+
|
| 36 |
+
# Needle Pi Coding Agent
|
| 37 |
+
|
| 38 |
+
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`.
|
| 39 |
+
|
| 40 |
+
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.
|
| 41 |
+
|
| 42 |
+
## Results
|
| 43 |
+
|
| 44 |
+
Evaluation used a held-out, balanced 40-example test set with 10 examples per tool.
|
| 45 |
+
|
| 46 |
+
| Metric | Base Needle | This fine-tune |
|
| 47 |
+
|---|---:|---:|
|
| 48 |
+
| Exact match | 27.5% | **80.0%** |
|
| 49 |
+
| Tool-call F1 | 28.95% | **80.0%** |
|
| 50 |
+
| Tool-name F1 | 88.0% | **100.0%** |
|
| 51 |
+
| JSON parse rate | 82.5% | **100.0%** |
|
| 52 |
+
| Argument accuracy | 32.35% | **80.0%** |
|
| 53 |
+
|
| 54 |
+
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`.
|
| 55 |
+
|
| 56 |
+
These are local single-run measurements, not an external benchmark. See [`metrics.json`](metrics.json) for machine-readable results.
|
| 57 |
+
|
| 58 |
+
## Training
|
| 59 |
+
|
| 60 |
+
- Base: `Cactus-Compute/needle`
|
| 61 |
+
- Examples: 1,500 (`bash`: 450; `read`, `edit`, `write`: 350 each)
|
| 62 |
+
- Split: 1,420 train / 40 validation / 40 test
|
| 63 |
+
- Epochs: 1
|
| 64 |
+
- Batch size: 32
|
| 65 |
+
- Training steps: 22
|
| 66 |
+
- Optimizer learning rates: Adam `3e-5`, Muon `0.02`
|
| 67 |
+
- Checkpoint-selection validation exact match: 70.0%
|
| 68 |
+
- Synthetic-data generator: Codex CLI, with strict schema, grounding, and deduplication validation
|
| 69 |
+
|
| 70 |
+
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.
|
| 71 |
+
|
| 72 |
+
## Usage
|
| 73 |
+
|
| 74 |
+
Clone the [Needle repository](https://github.com/cactus-compute/needle), install it, then download this repository:
|
| 75 |
+
|
| 76 |
+
```bash
|
| 77 |
+
hf download theabbie/needle-pi-coding-agent --local-dir needle-pi-coding-agent
|
| 78 |
+
```
|
| 79 |
+
|
| 80 |
+
```python
|
| 81 |
+
import json
|
| 82 |
+
from needle import SimpleAttentionNetwork, generate, get_tokenizer, load_checkpoint
|
| 83 |
+
|
| 84 |
+
params, config = load_checkpoint("needle-pi-coding-agent/needle-pi-coding-agent.pkl")
|
| 85 |
+
model = SimpleAttentionNetwork(config)
|
| 86 |
+
|
| 87 |
+
tools = json.loads(open("needle-pi-coding-agent/training/pi_tools.json").read())
|
| 88 |
+
result = generate(
|
| 89 |
+
model,
|
| 90 |
+
params,
|
| 91 |
+
get_tokenizer(),
|
| 92 |
+
query="Run bash command: curl -s https://api.ipify.org",
|
| 93 |
+
tools=json.dumps(tools, separators=(",", ":")),
|
| 94 |
+
stream=False,
|
| 95 |
+
)
|
| 96 |
+
print(result)
|
| 97 |
+
# [{"name":"bash","arguments":{"command":"curl -s https://api.ipify.org"}}]
|
| 98 |
+
```
|
| 99 |
+
|
| 100 |
+
The bundled [`integration/pi_tool_router.py`](integration/pi_tool_router.py) accepts this JSON on standard input:
|
| 101 |
+
|
| 102 |
+
```json
|
| 103 |
+
{"intent":"Run bash command: curl -s https://api.ipify.org","tools":[...]}
|
| 104 |
+
```
|
| 105 |
+
|
| 106 |
+
and emits a JSON array containing the selected call.
|
| 107 |
+
|
| 108 |
+
## Included artifacts
|
| 109 |
+
|
| 110 |
+
- `needle-pi-coding-agent.pkl`: best native Needle/JAX checkpoint
|
| 111 |
+
- `config.json` and `tokenizer/`: architecture and tokenizer files
|
| 112 |
+
- `training/pi_tools_1500.jsonl`: full 1,500-example training corpus
|
| 113 |
+
- `training/pi_tools.json`: flattened schemas for Pi's built-in tools
|
| 114 |
+
- `training/generate_pi_dataset.py`: resumable, atomically persisted generator
|
| 115 |
+
- `training/codex_client.py`: clean Codex CLI text-generation adapter
|
| 116 |
+
- `training/pi_tools_1500.progress.json`: final generation manifest
|
| 117 |
+
- `integration/pi_tool_router.py`: minimal stdin/stdout router used by the Pi provider prototype
|
| 118 |
+
|
| 119 |
+
## Limitations
|
| 120 |
+
|
| 121 |
+
- This is a focused single-shot tool router, not a conversational language model.
|
| 122 |
+
- Exact edit extraction is strongest when the planner explicitly supplies the `edits` array with `oldText` and `newText` values.
|
| 123 |
+
- Long or whitespace-sensitive write/edit payloads can still be copied imperfectly.
|
| 124 |
+
- Shell commands are executed by the surrounding agent runtime; this model does not provide sandboxing or command safety.
|
| 125 |
+
|
| 126 |
+
## License and attribution
|
| 127 |
+
|
| 128 |
+
MIT. Based on [Cactus-Compute/needle](https://huggingface.co/Cactus-Compute/needle) and the [Needle source repository](https://github.com/cactus-compute/needle).
|
config.json
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"library_name": "jax",
|
| 3 |
+
"model_type": "custom",
|
| 4 |
+
"architectures": ["SimpleAttentionNetwork"]
|
| 5 |
+
}
|
integration/pi_tool_router.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Route a short tool intent to Pi's live tool schemas with Needle."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import sys
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
from needle import SimpleAttentionNetwork, generate, get_tokenizer, load_checkpoint
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
MODEL_ROOT = Path(__file__).resolve().parents[1]
|
| 11 |
+
params, config = load_checkpoint(MODEL_ROOT / "needle-pi-coding-agent.pkl")
|
| 12 |
+
model = SimpleAttentionNetwork(config)
|
| 13 |
+
tokenizer = get_tokenizer()
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def main() -> None:
|
| 17 |
+
request = json.load(sys.stdin)
|
| 18 |
+
result = generate(
|
| 19 |
+
model,
|
| 20 |
+
params,
|
| 21 |
+
tokenizer,
|
| 22 |
+
query=request["intent"],
|
| 23 |
+
tools=json.dumps(request["tools"], separators=(",", ":")),
|
| 24 |
+
stream=False,
|
| 25 |
+
)
|
| 26 |
+
calls = json.loads(result)
|
| 27 |
+
if not isinstance(calls, list) or not calls:
|
| 28 |
+
raise ValueError("Needle returned no tool calls")
|
| 29 |
+
# Exa plans one step at a time. Running only Needle's first call keeps Pi's
|
| 30 |
+
# real tool result in the loop before either model chooses another action.
|
| 31 |
+
json.dump(calls[:1], sys.stdout, separators=(",", ":"))
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
if __name__ == "__main__":
|
| 35 |
+
main()
|
metrics.json
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"evaluation_split": {
|
| 3 |
+
"total": 40,
|
| 4 |
+
"per_tool": 10,
|
| 5 |
+
"tools": ["read", "bash", "edit", "write"]
|
| 6 |
+
},
|
| 7 |
+
"base_model": {
|
| 8 |
+
"call_f1": 0.2895,
|
| 9 |
+
"name_f1": 0.88,
|
| 10 |
+
"exact_match": 0.275,
|
| 11 |
+
"parse_rate": 0.825,
|
| 12 |
+
"arguments_accuracy": 0.3235,
|
| 13 |
+
"exact_by_tool": {"read": 0.8, "bash": 0.2, "edit": 0.0, "write": 0.1}
|
| 14 |
+
},
|
| 15 |
+
"fine_tuned_model": {
|
| 16 |
+
"call_f1": 0.8,
|
| 17 |
+
"name_f1": 1.0,
|
| 18 |
+
"exact_match": 0.8,
|
| 19 |
+
"parse_rate": 1.0,
|
| 20 |
+
"arguments_accuracy": 0.8,
|
| 21 |
+
"exact_by_tool": {"read": 0.9, "bash": 0.8, "edit": 0.9, "write": 0.6}
|
| 22 |
+
}
|
| 23 |
+
}
|
needle-pi-coding-agent.pkl
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:822e8749226c2289012adc778f2bea4ed288db630a2017ade301d730bd7d1cae
|
| 3 |
+
size 52632818
|
tokenizer/needle.model
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0823f5b9133c68a8140addc5d7a425fa9119c4c8cb4a550363b4bffa4ba1c8c7
|
| 3 |
+
size 124960
|
tokenizer/needle.vocab
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
training/codex_client.py
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Small text-generation adapter backed by `codex exec`."""
|
| 2 |
+
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
from types import SimpleNamespace
|
| 5 |
+
import subprocess
|
| 6 |
+
import tempfile
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class CodexModels:
|
| 10 |
+
def generate_content(self, model, contents, config=None):
|
| 11 |
+
del model, config
|
| 12 |
+
with tempfile.TemporaryDirectory(prefix="needle-codex-") as temp_dir:
|
| 13 |
+
output_path = Path(temp_dir) / "response.txt"
|
| 14 |
+
process = subprocess.run(
|
| 15 |
+
[
|
| 16 |
+
"codex",
|
| 17 |
+
"exec",
|
| 18 |
+
"--ephemeral",
|
| 19 |
+
"--sandbox",
|
| 20 |
+
"read-only",
|
| 21 |
+
"--skip-git-repo-check",
|
| 22 |
+
"--color",
|
| 23 |
+
"never",
|
| 24 |
+
"--output-last-message",
|
| 25 |
+
str(output_path),
|
| 26 |
+
"-",
|
| 27 |
+
],
|
| 28 |
+
input=contents,
|
| 29 |
+
text=True,
|
| 30 |
+
stdout=subprocess.PIPE,
|
| 31 |
+
stderr=subprocess.PIPE,
|
| 32 |
+
timeout=600,
|
| 33 |
+
check=False,
|
| 34 |
+
)
|
| 35 |
+
if process.returncode != 0:
|
| 36 |
+
detail = (process.stderr or process.stdout).strip()[-4000:]
|
| 37 |
+
raise RuntimeError(f"codex exec failed ({process.returncode}): {detail}")
|
| 38 |
+
if not output_path.exists():
|
| 39 |
+
raise RuntimeError("codex exec produced no final-response file")
|
| 40 |
+
text = output_path.read_text(encoding="utf-8").strip()
|
| 41 |
+
if not text:
|
| 42 |
+
raise RuntimeError("codex exec produced an empty final response")
|
| 43 |
+
return SimpleNamespace(text=text)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
class CodexClient:
|
| 47 |
+
def __init__(self):
|
| 48 |
+
self.models = CodexModels()
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
class CodexClientPool:
|
| 52 |
+
"""Matches the tiny client-pool interface used by Needle's generator."""
|
| 53 |
+
|
| 54 |
+
def __init__(self):
|
| 55 |
+
self._client = CodexClient()
|
| 56 |
+
|
| 57 |
+
def get(self):
|
| 58 |
+
return self._client
|
training/generate_pi_dataset.py
ADDED
|
@@ -0,0 +1,291 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Resumably extend Pi's tool-routing dataset with harder parser examples."""
|
| 2 |
+
|
| 3 |
+
from collections import Counter
|
| 4 |
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import Any
|
| 7 |
+
import json
|
| 8 |
+
import os
|
| 9 |
+
import re
|
| 10 |
+
import time
|
| 11 |
+
|
| 12 |
+
from codex_client import CodexClientPool
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
ROOT = Path(__file__).parent
|
| 16 |
+
TOOLS = json.loads((ROOT / "pi_tools.json").read_text(encoding="utf-8"))
|
| 17 |
+
SEED_OUTPUT = ROOT / "pi_tools_1200.jsonl"
|
| 18 |
+
OUTPUT = ROOT / "pi_tools_1500.jsonl"
|
| 19 |
+
PROGRESS = ROOT / "pi_tools_1500.progress.json"
|
| 20 |
+
TARGETS = {"read": 350, "bash": 450, "edit": 350, "write": 350}
|
| 21 |
+
TOTAL_TARGET = sum(TARGETS.values())
|
| 22 |
+
BATCH_SIZE = 25
|
| 23 |
+
MAX_WORKERS = 4
|
| 24 |
+
TOOL_MAP = {tool["name"]: tool for tool in TOOLS}
|
| 25 |
+
TOOLS_TEXT = json.dumps(TOOLS, separators=(",", ":"), ensure_ascii=False)
|
| 26 |
+
|
| 27 |
+
THEMES = [
|
| 28 |
+
"short direct requests",
|
| 29 |
+
"absolute and relative paths",
|
| 30 |
+
"spaces, quotes, unicode, and punctuation",
|
| 31 |
+
"JavaScript and TypeScript projects",
|
| 32 |
+
"Python projects",
|
| 33 |
+
"shell pipelines and redirects",
|
| 34 |
+
"configuration and environment files",
|
| 35 |
+
"multiline source code",
|
| 36 |
+
"logs, offsets, and limits",
|
| 37 |
+
"nested directories and dotfiles",
|
| 38 |
+
"URLs contrasted with local paths",
|
| 39 |
+
"realistic debugging and maintenance tasks",
|
| 40 |
+
]
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def _strip_fence(text: str) -> str:
|
| 44 |
+
text = text.strip()
|
| 45 |
+
if text.startswith("```"):
|
| 46 |
+
text = text.split("\n", 1)[1] if "\n" in text else text[3:]
|
| 47 |
+
if text.endswith("```"):
|
| 48 |
+
text = text[:-3]
|
| 49 |
+
return text.strip()
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _query_key(query: str) -> str:
|
| 53 |
+
return re.sub(r"\s+", " ", query.strip().casefold())
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _load_existing() -> list[dict[str, Any]]:
|
| 57 |
+
source = OUTPUT if OUTPUT.exists() else SEED_OUTPUT
|
| 58 |
+
if not source.exists():
|
| 59 |
+
return []
|
| 60 |
+
records = []
|
| 61 |
+
for line_number, line in enumerate(source.read_text(encoding="utf-8").splitlines(), 1):
|
| 62 |
+
if not line.strip():
|
| 63 |
+
continue
|
| 64 |
+
try:
|
| 65 |
+
record = json.loads(line)
|
| 66 |
+
answers = json.loads(record["answers"])
|
| 67 |
+
name = answers[0]["name"]
|
| 68 |
+
if name not in TOOL_MAP:
|
| 69 |
+
raise ValueError(f"unknown tool {name}")
|
| 70 |
+
records.append(record)
|
| 71 |
+
except Exception as error:
|
| 72 |
+
raise ValueError(f"invalid existing dataset line {line_number}: {error}") from error
|
| 73 |
+
return records
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def _leaf_values(value: Any):
|
| 77 |
+
if isinstance(value, dict):
|
| 78 |
+
for child in value.values():
|
| 79 |
+
yield from _leaf_values(child)
|
| 80 |
+
elif isinstance(value, list):
|
| 81 |
+
for child in value:
|
| 82 |
+
yield from _leaf_values(child)
|
| 83 |
+
elif value is not None:
|
| 84 |
+
yield value
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def _valid_example(example: Any, expected_tool: str) -> tuple[bool, str]:
|
| 88 |
+
if not isinstance(example, dict):
|
| 89 |
+
return False, "not an object"
|
| 90 |
+
query = example.get("query")
|
| 91 |
+
answers = example.get("answers")
|
| 92 |
+
if not isinstance(query, str) or not query.strip():
|
| 93 |
+
return False, "missing query"
|
| 94 |
+
if not isinstance(answers, list) or len(answers) != 1:
|
| 95 |
+
return False, "answer count"
|
| 96 |
+
call = answers[0]
|
| 97 |
+
if not isinstance(call, dict) or call.get("name") != expected_tool:
|
| 98 |
+
return False, "wrong tool"
|
| 99 |
+
arguments = call.get("arguments")
|
| 100 |
+
if not isinstance(arguments, dict):
|
| 101 |
+
return False, "arguments"
|
| 102 |
+
|
| 103 |
+
schema = TOOL_MAP[expected_tool]["parameters"]
|
| 104 |
+
if set(arguments) - set(schema):
|
| 105 |
+
return False, "unknown argument"
|
| 106 |
+
required = {name for name, info in schema.items() if info.get("required")}
|
| 107 |
+
if required - set(arguments):
|
| 108 |
+
return False, "missing argument"
|
| 109 |
+
|
| 110 |
+
for name, value in arguments.items():
|
| 111 |
+
expected_type = schema[name]["type"]
|
| 112 |
+
if expected_type == "string" and not isinstance(value, str):
|
| 113 |
+
return False, "string type"
|
| 114 |
+
if expected_type == "number" and (not isinstance(value, (int, float)) or isinstance(value, bool)):
|
| 115 |
+
return False, "number type"
|
| 116 |
+
if expected_type == "array" and not isinstance(value, list):
|
| 117 |
+
return False, "array type"
|
| 118 |
+
|
| 119 |
+
if expected_tool == "edit":
|
| 120 |
+
edits = arguments.get("edits")
|
| 121 |
+
if not edits or not all(
|
| 122 |
+
isinstance(edit, dict)
|
| 123 |
+
and set(edit) == {"oldText", "newText"}
|
| 124 |
+
and isinstance(edit["oldText"], str)
|
| 125 |
+
and isinstance(edit["newText"], str)
|
| 126 |
+
for edit in edits
|
| 127 |
+
):
|
| 128 |
+
return False, "edit shape"
|
| 129 |
+
|
| 130 |
+
# The router must copy values, never invent them. Require every non-empty
|
| 131 |
+
# scalar argument value to occur literally in the planner-style query.
|
| 132 |
+
for value in _leaf_values(arguments):
|
| 133 |
+
rendered = str(value)
|
| 134 |
+
if rendered and rendered not in query:
|
| 135 |
+
return False, f"ungrounded value: {rendered[:40]}"
|
| 136 |
+
return True, ""
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def _prompt(tool_name: str, count: int, batch_number: int, avoid: list[str], current_count: int) -> str:
|
| 140 |
+
theme = THEMES[(batch_number - 1) % len(THEMES)]
|
| 141 |
+
if tool_name == "edit":
|
| 142 |
+
theme = "short single-line exact replacements with diverse file types and syntax"
|
| 143 |
+
if tool_name == "bash" and current_count < 400:
|
| 144 |
+
theme = """high-complexity shell-command parsing. Invent realistic setup or maintenance scenarios such as installing X, configuring Y, launching a service, building a project, transforming files, inspecting a system, or chaining arbitrary shell steps. The query must already contain the exact complete command; the router only extracts it. Use varied quoting, flags, environment variables, pipes, redirects, subshells, and multiline commands. Do not require the router to synthesize specialist commands from vague goals"""
|
| 145 |
+
elif current_count >= 300:
|
| 146 |
+
theme = """noisy or mildly ambiguous planner messages that still contain one unambiguous tool action and every exact argument. Add natural hesitations, corrections, irrelevant context, or conversational clutter, but keep the intended call objectively recoverable"""
|
| 147 |
+
avoid_text = "\n".join(f"- {query}" for query in avoid[-30:]) or "- none"
|
| 148 |
+
tool = TOOL_MAP[tool_name]
|
| 149 |
+
return f"""Generate exactly {count} training examples for a coding-agent tool router.
|
| 150 |
+
|
| 151 |
+
ALL AVAILABLE TOOLS:
|
| 152 |
+
{TOOLS_TEXT}
|
| 153 |
+
|
| 154 |
+
This batch must call only `{tool_name}`:
|
| 155 |
+
{json.dumps(tool, separators=(',', ':'), ensure_ascii=False)}
|
| 156 |
+
|
| 157 |
+
Batch theme: {theme}
|
| 158 |
+
|
| 159 |
+
Rules:
|
| 160 |
+
- Return one valid JSON array and nothing else.
|
| 161 |
+
- Every item is exactly {{"query":"...","answers":[{{"name":"{tool_name}","arguments":{{...}}}}]}}.
|
| 162 |
+
- Exactly one call per example and exactly {count} examples.
|
| 163 |
+
- The query is a short, explicit planner intent, not a casual end-user request.
|
| 164 |
+
- Include every argument value literally and exactly in the query. Never infer, shorten, repair, normalize, or invent a value.
|
| 165 |
+
- Preserve shell commands, paths, file content, oldText, newText, quoting, escapes, whitespace, and newlines exactly.
|
| 166 |
+
- Use only declared parameters and include every required parameter.
|
| 167 |
+
- Make examples meaningfully different, not template-like substitutions.
|
| 168 |
+
- For `bash`, `command` must contain the complete executable command, not merely a URL, filename, or description.
|
| 169 |
+
- For `read`, paths are filesystem paths; never use a web URL as a path.
|
| 170 |
+
- For `edit`, `edits` is a non-empty array of objects with only `oldText` and `newText` strings.
|
| 171 |
+
- For `edit`, keep every oldText and newText on one line, and place each exact value directly in the query without Markdown escaping or a code fence.
|
| 172 |
+
- For `write`, `content` is the complete file content and appears verbatim in the query.
|
| 173 |
+
- Noise may surround an action, but never make the expected arguments unknowable or require specialist domain knowledge.
|
| 174 |
+
|
| 175 |
+
Do not repeat these earlier queries:
|
| 176 |
+
{avoid_text}
|
| 177 |
+
"""
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
def _generate_batch(tool_name: str, count: int, batch_number: int, avoid: list[str], current_count: int):
|
| 181 |
+
response = CodexClientPool().get().models.generate_content(
|
| 182 |
+
model="codex-cli",
|
| 183 |
+
contents=_prompt(tool_name, count, batch_number, avoid, current_count),
|
| 184 |
+
config={"temperature": 0.8, "max_output_tokens": 16384},
|
| 185 |
+
)
|
| 186 |
+
parsed = json.loads(_strip_fence(response.text))
|
| 187 |
+
if not isinstance(parsed, list):
|
| 188 |
+
raise ValueError("response is not an array")
|
| 189 |
+
return tool_name, parsed
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def _persist(records: list[dict[str, Any]], failures: dict[str, int]) -> Counter:
|
| 193 |
+
counts = Counter()
|
| 194 |
+
lines = []
|
| 195 |
+
for record in records:
|
| 196 |
+
name = json.loads(record["answers"])[0]["name"]
|
| 197 |
+
counts[name] += 1
|
| 198 |
+
lines.append(json.dumps(record, ensure_ascii=False))
|
| 199 |
+
|
| 200 |
+
output_tmp = OUTPUT.with_suffix(OUTPUT.suffix + ".tmp")
|
| 201 |
+
output_tmp.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
|
| 202 |
+
os.replace(output_tmp, OUTPUT)
|
| 203 |
+
|
| 204 |
+
state = {
|
| 205 |
+
"targets": TARGETS,
|
| 206 |
+
"total": len(records),
|
| 207 |
+
"per_tool": {name: counts[name] for name in TOOL_MAP},
|
| 208 |
+
"failures": failures,
|
| 209 |
+
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
| 210 |
+
"complete": all(counts[name] >= TARGETS[name] for name in TOOL_MAP),
|
| 211 |
+
}
|
| 212 |
+
progress_tmp = PROGRESS.with_suffix(PROGRESS.suffix + ".tmp")
|
| 213 |
+
progress_tmp.write_text(json.dumps(state, indent=2) + "\n", encoding="utf-8")
|
| 214 |
+
os.replace(progress_tmp, PROGRESS)
|
| 215 |
+
return counts
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def main() -> None:
|
| 219 |
+
records = _load_existing()
|
| 220 |
+
seen = {_query_key(record["query"]) for record in records}
|
| 221 |
+
failures = {name: 0 for name in TOOL_MAP}
|
| 222 |
+
counts = _persist(records, failures)
|
| 223 |
+
print(f"Resuming at {len(records)}/{TOTAL_TARGET}: {dict(counts)}", flush=True)
|
| 224 |
+
|
| 225 |
+
round_number = 0
|
| 226 |
+
while any(counts[name] < TARGETS[name] for name in TOOL_MAP):
|
| 227 |
+
round_number += 1
|
| 228 |
+
jobs = {}
|
| 229 |
+
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
|
| 230 |
+
for name in TOOL_MAP:
|
| 231 |
+
remaining = TARGETS[name] - counts[name]
|
| 232 |
+
if remaining <= 0:
|
| 233 |
+
continue
|
| 234 |
+
count = min(BATCH_SIZE, remaining)
|
| 235 |
+
prior = [
|
| 236 |
+
record["query"] for record in records
|
| 237 |
+
if json.loads(record["answers"])[0]["name"] == name
|
| 238 |
+
]
|
| 239 |
+
batch_number = (counts[name] // BATCH_SIZE) + 1
|
| 240 |
+
future = pool.submit(_generate_batch, name, count, batch_number, prior, counts[name])
|
| 241 |
+
jobs[future] = name
|
| 242 |
+
|
| 243 |
+
for future in as_completed(jobs):
|
| 244 |
+
scheduled_name = jobs[future]
|
| 245 |
+
try:
|
| 246 |
+
name, examples = future.result()
|
| 247 |
+
accepted = 0
|
| 248 |
+
rejected = Counter()
|
| 249 |
+
remaining = TARGETS[name] - counts[name]
|
| 250 |
+
for example in examples:
|
| 251 |
+
valid, reason = _valid_example(example, name)
|
| 252 |
+
key = _query_key(example.get("query", "")) if isinstance(example, dict) else ""
|
| 253 |
+
if not valid:
|
| 254 |
+
rejected[reason] += 1
|
| 255 |
+
continue
|
| 256 |
+
if key in seen:
|
| 257 |
+
rejected["duplicate"] += 1
|
| 258 |
+
continue
|
| 259 |
+
if accepted >= remaining:
|
| 260 |
+
break
|
| 261 |
+
answers = example["answers"]
|
| 262 |
+
records.append({
|
| 263 |
+
"query": example["query"],
|
| 264 |
+
"tools": TOOLS_TEXT,
|
| 265 |
+
"answers": json.dumps(answers, separators=(",", ":"), ensure_ascii=False),
|
| 266 |
+
"source": "synth-codex-cli",
|
| 267 |
+
"model": "codex-cli",
|
| 268 |
+
})
|
| 269 |
+
seen.add(key)
|
| 270 |
+
accepted += 1
|
| 271 |
+
failures[name] = 0 if accepted else failures[name] + 1
|
| 272 |
+
counts = _persist(records, failures)
|
| 273 |
+
print(
|
| 274 |
+
f"{name}: +{accepted}, rejected={dict(rejected)}, "
|
| 275 |
+
f"progress={counts[name]}/{TARGETS[name]}, total={len(records)}/{TOTAL_TARGET}",
|
| 276 |
+
flush=True,
|
| 277 |
+
)
|
| 278 |
+
except Exception as error:
|
| 279 |
+
failures[scheduled_name] += 1
|
| 280 |
+
_persist(records, failures)
|
| 281 |
+
print(f"{scheduled_name} batch failed (progress preserved): {error}", flush=True)
|
| 282 |
+
|
| 283 |
+
if max(failures.values()) >= 5:
|
| 284 |
+
raise RuntimeError(f"five consecutive failures; rerun to resume: {failures}")
|
| 285 |
+
|
| 286 |
+
counts = _persist(records, failures)
|
| 287 |
+
print(json.dumps({"output": str(OUTPUT), "total": len(records), "per_tool": dict(counts)}, indent=2))
|
| 288 |
+
|
| 289 |
+
|
| 290 |
+
if __name__ == "__main__":
|
| 291 |
+
main()
|
training/pi_tools.json
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"name": "read",
|
| 4 |
+
"description": "Read the contents of a file. Supports text files and images. Use offset and limit for large text files.",
|
| 5 |
+
"parameters": {
|
| 6 |
+
"path": {"type": "string", "description": "Path to the file to read (relative or absolute)", "required": true},
|
| 7 |
+
"offset": {"type": "number", "description": "Line number to start reading from (1-indexed)", "required": false},
|
| 8 |
+
"limit": {"type": "number", "description": "Maximum number of lines to read", "required": false}
|
| 9 |
+
}
|
| 10 |
+
},
|
| 11 |
+
{
|
| 12 |
+
"name": "bash",
|
| 13 |
+
"description": "Execute a bash command in the current working directory and return stdout and stderr.",
|
| 14 |
+
"parameters": {
|
| 15 |
+
"command": {"type": "string", "description": "Exact Bash command to execute", "required": true},
|
| 16 |
+
"timeout": {"type": "number", "description": "Optional timeout in seconds", "required": false}
|
| 17 |
+
}
|
| 18 |
+
},
|
| 19 |
+
{
|
| 20 |
+
"name": "edit",
|
| 21 |
+
"description": "Edit one file using one or more exact, non-overlapping text replacements.",
|
| 22 |
+
"parameters": {
|
| 23 |
+
"path": {"type": "string", "description": "Path to the file to edit (relative or absolute)", "required": true},
|
| 24 |
+
"edits": {"type": "array", "description": "Array of objects with exact oldText and replacement newText strings", "required": true}
|
| 25 |
+
}
|
| 26 |
+
},
|
| 27 |
+
{
|
| 28 |
+
"name": "write",
|
| 29 |
+
"description": "Create or completely overwrite a file, creating parent directories automatically.",
|
| 30 |
+
"parameters": {
|
| 31 |
+
"path": {"type": "string", "description": "Path to the file to write (relative or absolute)", "required": true},
|
| 32 |
+
"content": {"type": "string", "description": "Complete content to write to the file", "required": true}
|
| 33 |
+
}
|
| 34 |
+
}
|
| 35 |
+
]
|
training/pi_tools_1500.jsonl
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
training/pi_tools_1500.progress.json
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"targets": {
|
| 3 |
+
"read": 350,
|
| 4 |
+
"bash": 450,
|
| 5 |
+
"edit": 350,
|
| 6 |
+
"write": 350
|
| 7 |
+
},
|
| 8 |
+
"total": 1500,
|
| 9 |
+
"per_tool": {
|
| 10 |
+
"read": 350,
|
| 11 |
+
"bash": 450,
|
| 12 |
+
"edit": 350,
|
| 13 |
+
"write": 350
|
| 14 |
+
},
|
| 15 |
+
"failures": {
|
| 16 |
+
"read": 0,
|
| 17 |
+
"bash": 0,
|
| 18 |
+
"edit": 0,
|
| 19 |
+
"write": 0
|
| 20 |
+
},
|
| 21 |
+
"updated_at": "2026-07-13T11:05:43+0530",
|
| 22 |
+
"complete": true
|
| 23 |
+
}
|