File size: 2,374 Bytes
c128a6c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """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): # the connector writes states asynchronously
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)
|