#!/usr/bin/env python3
"""SC-6: held-out perplexity via the completions API (echo+logprobs).
Usage: sc6_ppl.py --port 8199 [--ref /data/glm52-heldout.txt]"""
import argparse, json, math, urllib.request

ap = argparse.ArgumentParser()
ap.add_argument("--port", type=int, required=True)
ap.add_argument("--ref", default="/data/glm52-heldout.txt")
ap.add_argument("--chunk-tokens", type=int, default=3500)
a = ap.parse_args()

model = json.load(urllib.request.urlopen(
    f"http://localhost:{a.port}/v1/models"))["data"][0]["id"]
text = open(a.ref, errors="ignore").read()
# split by chars ~4/token into chunks
step = a.chunk_tokens * 4
tot_lp, tot_n = 0.0, 0
for i in range(0, len(text), step):
    chunk = text[i:i + step]
    if len(chunk) < 2000: break
    body = {"model": model, "prompt": chunk, "max_tokens": 0,
            "echo": True, "logprobs": 0}
    req = urllib.request.Request(
        f"http://localhost:{a.port}/v1/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=600) as r:
        resp = json.load(r)
    lps = resp["choices"][0]["logprobs"]["token_logprobs"]
    lps = [x for x in lps if x is not None]
    tot_lp += sum(lps); tot_n += len(lps)
ppl = math.exp(-tot_lp / max(tot_n, 1))
print(f"SC6 ppl={ppl:.4f} over {tot_n} tokens")