| """Generates completions for eval/train JSONL files through a running llama-server.
|
|
|
| Prompts are rendered with the model's own HF chat template (thinking enabled) and sent as token ids, so the
|
| server sees exactly the same input for every model. Output is appended to a JSONL file and the run resumes
|
| where it stopped.
|
|
|
| usage: gen.py --inp data/eval/gsm8k.jsonl --out results/orig/gsm8k.jsonl [--n 1] [--max-tokens 16384]
|
| """
|
| import argparse
|
| import json
|
| import os
|
| import threading
|
| import time
|
| import zlib
|
| import urllib.request
|
| from concurrent.futures import ThreadPoolExecutor
|
|
|
| from transformers import AutoTokenizer
|
|
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--inp", nargs="+", required=True)
|
| ap.add_argument("--out", required=True)
|
| ap.add_argument("--n", type=int, default=1, help="samples per prompt")
|
| ap.add_argument("--max-tokens", type=int, default=16384)
|
| ap.add_argument("--temperature", type=float, default=0.6)
|
| ap.add_argument("--top-p", type=float, default=0.95)
|
| ap.add_argument("--top-k", type=int, default=50)
|
| ap.add_argument("--save-tokens", action="store_true", help="store generated token ids (for RL training)")
|
| ap.add_argument("--workers", type=int, default=32)
|
| ap.add_argument("--limit", type=int, default=0)
|
| ap.add_argument("--port", type=int, default=8080)
|
| ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
|
| ap.add_argument("--seed", type=int, default=0)
|
| ap.add_argument("--system-file", default="", help="optional system prompt (text file); json {task: prompt} also works")
|
| args = ap.parse_args()
|
|
|
| tok = AutoTokenizer.from_pretrained(args.tokenizer)
|
| SYSTEM = ""
|
| if args.system_file:
|
| SYSTEM = open(args.system_file, encoding="utf-8").read()
|
| if args.system_file.endswith(".json"):
|
| SYSTEM = json.loads(SYSTEM)
|
| items = []
|
| for path in args.inp:
|
| with open(path, encoding="utf-8") as f:
|
| rows = [json.loads(l) for l in f]
|
| items += rows[: args.limit] if args.limit else rows
|
|
|
| done = set()
|
| if os.path.exists(args.out):
|
| with open(args.out, encoding="utf-8") as f:
|
| for l in f:
|
| r = json.loads(l)
|
| done.add((r["id"], r["sample"]))
|
| os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
|
|
|
| jobs = [(it, s) for it in items for s in range(args.n) if (it["id"], s) not in done]
|
| print(f"{len(jobs)} jobs ({len(done)} already done)", flush=True)
|
|
|
| lock = threading.Lock()
|
| out_f = open(args.out, "a", encoding="utf-8")
|
| stats = {"n": 0, "tok": 0, "t0": time.time()}
|
|
|
|
|
| def run(job):
|
| it, s = job
|
| sysp = SYSTEM.get(it["task"], SYSTEM.get("*", "")) if isinstance(SYSTEM, dict) else SYSTEM
|
| ids = tok.apply_chat_template(([{"role": "system", "content": sysp}] if sysp else []) +
|
| [{"role": "user", "content": it["prompt"]}], tokenize=True,
|
| add_generation_prompt=True, return_dict=False)
|
| if isinstance(ids, dict):
|
| ids = ids["input_ids"]
|
| body = {
|
| "prompt": ids, "n_predict": args.max_tokens, "temperature": args.temperature, "top_p": args.top_p,
|
| "top_k": args.top_k, "min_p": 0.0, "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0,
|
| "dry_multiplier": 0.0, "xtc_probability": 0.0, "typical_p": 1.0,
|
| "seed": args.seed * 1000003 + zlib.crc32(it["id"].encode()) % 1000003 * 7 + s, "cache_prompt": False,
|
| "samplers": ["top_k", "top_p", "temperature"], "return_tokens": args.save_tokens,
|
| }
|
| for attempt in range(5):
|
| try:
|
| req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
|
| headers={"Content-Type": "application/json"})
|
| res = json.loads(urllib.request.urlopen(req, timeout=3600).read())
|
| break
|
| except Exception as e:
|
| if attempt == 4:
|
| raise
|
| time.sleep(5)
|
| rec = {"id": it["id"], "task": it["task"], "sample": s, "completion": "<think>" + res["content"],
|
| "n_tokens": res["tokens_predicted"], "truncated": res.get("stop_type") == "limit",
|
| "prompt_tokens": len(ids), "stop_type": res.get("stop_type")}
|
| if args.save_tokens:
|
| rec["tokens"] = res["tokens"]
|
| with lock:
|
| out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
| out_f.flush()
|
| stats["n"] += 1
|
| stats["tok"] += rec["n_tokens"]
|
| if stats["n"] % 50 == 0 or stats["n"] == len(jobs):
|
| dt = time.time() - stats["t0"]
|
| print(f"{stats['n']}/{len(jobs)} {stats['tok'] / dt:.0f} tok/s avg_len {stats['tok'] / stats['n']:.0f} "
|
| f"elapsed {dt / 60:.1f} min", flush=True)
|
|
|
|
|
| with ThreadPoolExecutor(args.workers) as ex:
|
| list(ex.map(run, jobs))
|
| out_f.close()
|
|
|
