56 lines
2.2 KiB
Python
Executable File
56 lines
2.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""deep-probe.py — one chat completion at ~43.6k prompt tokens, temp 0, 512
|
|
generated, streamed; decode = completion_tokens / (total - ttft)."""
|
|
import json, time, urllib.request
|
|
|
|
TARGET_BYTES = int(43600 * 3.8)
|
|
|
|
def build():
|
|
parts, n = [], 0
|
|
with open("/home/dank/corpora/out/calib.jsonl") as f:
|
|
for line in f:
|
|
p = json.loads(line)["prompt"]
|
|
parts.append(p)
|
|
n += len(p.encode())
|
|
if n >= TARGET_BYTES:
|
|
break
|
|
return "".join(parts)
|
|
|
|
def main(salt):
|
|
body = json.dumps({
|
|
"model": "frank-primary",
|
|
"messages": [{"role": "user", "content": f"[probe {salt}] Continue and summarize this material in English, 400 words.\n\n" + build()}],
|
|
"temperature": 0.0, "max_tokens": 512,
|
|
"stream": True, "stream_options": {"include_usage": True},
|
|
}).encode()
|
|
req = urllib.request.Request("http://localhost:7770/v1/chat/completions",
|
|
data=body, headers={"Content-Type": "application/json"})
|
|
t0 = time.time(); ttft = None; ctok = 0; usage = {}
|
|
with urllib.request.urlopen(req, timeout=2400) as r:
|
|
for raw in r:
|
|
line = raw.decode().strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
d = line[5:].strip()
|
|
if d == "[DONE]":
|
|
break
|
|
j = json.loads(d)
|
|
if j.get("usage"):
|
|
usage = j["usage"]
|
|
for ch in j.get("choices", []):
|
|
delta = (ch.get("delta") or {})
|
|
piece = delta.get("content") or delta.get("reasoning_content") or ""
|
|
if piece and ttft is None:
|
|
ttft = time.time() - t0
|
|
ctok += 1
|
|
total = time.time() - t0
|
|
gen = usage.get("completion_tokens") or ctok
|
|
dec = gen / (total - (ttft or 0)) if total > (ttft or 0) else 0
|
|
print(json.dumps({"salt": salt, "prompt_tokens": usage.get("prompt_tokens"),
|
|
"completion_tokens": gen, "ttft_s": round(ttft or 0, 2),
|
|
"total_s": round(total, 1), "decode_tps": round(dec, 2)}))
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
main(sys.argv[1] if len(sys.argv) > 1 else "x")
|