File size: 4,874 Bytes
67ebcef | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """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: # server busy / transient
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()
|