| """Generate teacher trajectories for training samples: prompts -> teacher answers + states files, |
| kept in /root/states. The trainer reads these. |
| |
| .venv/bin/python gen_teacher.py --dataset drawn.jsonl --samples 32 --max-prompt 1500 |
| """ |
| import argparse |
| import json |
| import time |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| from eval_service import dataset, run, servers |
| from eval_service.run import _render, generate, load_tokenizer |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--dataset", type=Path, default=Path("/root/drawn.jsonl")) |
| ap.add_argument("--teacher", type=Path, default=Path("/root/models/teacher")) |
| ap.add_argument("--samples", type=int, default=32) |
| ap.add_argument("--max-prompt", type=int, default=1500) |
| ap.add_argument("--workers", type=int, default=8) |
| ap.add_argument("--out", type=Path, default=Path("/root/train_samples.json")) |
| args = ap.parse_args() |
|
|
| samples = dataset.load(args.dataset).samples |
| tokenizer = load_tokenizer(args.teacher) |
| picked, prompts = [], [] |
| for s in samples: |
| p = _render(tokenizer, s.messages, s.tools) |
| if len(p) <= args.max_prompt: |
| picked.append(s) |
| prompts.append(p) |
| if len(picked) >= args.samples: |
| break |
| print(f"{len(picked)} samples, prompts {min(map(len, prompts))}-{max(map(len, prompts))} tokens") |
|
|
| out = Path("/root/states") |
| out.mkdir(exist_ok=True) |
| records = [] |
|
|
|
|
| def one(i, s, p): |
| answer, tpath = generate(groups["teacher"], p, i, s.id) |
| for _ in range(120): |
| if Path(tpath).exists(): |
| break |
| time.sleep(0.5) |
| keep = out / f"t{i}.safetensors" |
| keep.write_bytes(Path(tpath).read_bytes()) |
| Path(str(keep) + ".lock").touch() |
| return answer, str(keep) |
|
|
|
|
| try: |
| groups = servers.serve({"teacher": args.teacher}) |
| with ThreadPoolExecutor(args.workers) as pool: |
| futures = {i: pool.submit(one, i, s, p) for i, (s, p) in enumerate(zip(picked, prompts))} |
| for i, fut in futures.items(): |
| answer, keep = fut.result() |
| records.append({"i": i, "id": picked[i].id, "prompt": prompts[i], "answer": answer, "states": keep}) |
| print(f"{i+1}/{len(picked)} {picked[i].id} p{len(prompts[i])} a{len(answer)}", flush=True) |
| finally: |
| servers.stop_all() |
|
|
| args.out.write_text(json.dumps(records)) |
| print("saved", args.out) |
|
|
