# E20 generator: sends prompts to a running llama-server and writes one jsonl line per sample.
# Resumable: skips (item, sample) pairs already in --out.
import json, time, argparse, requests
from concurrent.futures import ThreadPoolExecutor
SYS = "Please reason step by step, and put your final answer within \\boxed{}."
def prompt(q):  # Qwen3 chat template with thinking disabled (enable_thinking=False)
    return ("<|im_start|>system\n" + SYS + "<|im_end|>\n<|im_start|>user\n" + q +
            "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n")
ap = argparse.ArgumentParser()
ap.add_argument("--items", required=True)      # jsonl with item, problem, gold, level, subject
ap.add_argument("--out", required=True)
ap.add_argument("--url", default="http://127.0.0.1:8096")
ap.add_argument("--n", type=int, default=1)    # samples per item
ap.add_argument("--temp", type=float, default=0.0)
ap.add_argument("--cap", type=int, default=2048)
ap.add_argument("--par", type=int, default=8)
a = ap.parse_args()
items = [json.loads(l) for l in open(a.items, encoding="utf-8")]
done = set()
try:
    for l in open(a.out, encoding="utf-8"):
        r = json.loads(l); done.add((r["item"], r["sample"]))
except FileNotFoundError: pass
jobs = [(it, s) for it in items for s in range(a.n) if (it["item"], s) not in done]
print("todo", len(jobs), flush=True)
def run(job):
    it, s = job
    body = {"prompt": prompt(it["problem"]), "n_predict": a.cap, "cache_prompt": False, "seed": 1000 * s + 7}
    if a.temp == 0: body.update(temperature=0, top_k=1)
    else: body.update(temperature=a.temp, top_p=0.95, top_k=40)
    t = time.time()
    j = requests.post(a.url + "/completion", json=body, timeout=3600).json()
    tm = j.get("timings", {})
    return dict(it, sample=s, text=j.get("content", ""), predicted_n=tm.get("predicted_n"),
                prompt_n=tm.get("prompt_n"), wall=round(time.time() - t, 2))
t0 = time.time(); k = 0
with open(a.out, "a", encoding="utf-8") as f, ThreadPoolExecutor(a.par) as ex:
    for r in ex.map(run, jobs):
        f.write(json.dumps(r) + "\n"); f.flush(); k += 1
        if k % 25 == 0: print(k, len(jobs), round(time.time() - t0), flush=True)
print("GENDONE", round(time.time() - t0), flush=True)