test-time-compute / tools /extend_sequential.py
ChandraPrakashBathula's picture
Space deploy: full app minus figures/ (not needed by the running Docker app; HF's binary-content check rejected those PNGs)
c140d3a
Raw History Blame Contribute Delete
5.77 kB
#!/usr/bin/env python3
"""Extend the sequential chains of a completed run so budgets can be compared.
The main experiment stops at K=4, which is enough to see whether revision helps
but not enough to answer the question the field actually cares about:
given a fixed inference budget, is it better spent going wide or going deep?
Comparing N=16 against K=4 does not answer that, because they are not the same
budget. This continues each recorded chain to K=16 so parallel and sequential can
be read at matched budgets B in {1,2,4,8,16}.
The chains are *continued*, not regenerated. Each stored round has its response
text, so the exact conversation the original run built is reconstructed and
extended, which keeps rounds 1-4 identical to what the main experiment already
reported rather than re-rolling them.
Also records per-call latency, which the main run does not. Tokens measure what
you pay; latency measures what you wait, and the two arms differ on the second in
a way tokens alone hide: 16 samples can be issued at once, while 16 revisions
cannot start until the previous one finishes.
python3 tools/extend_sequential.py <cache_key> [--rounds=16]
"""
from __future__ import annotations
import asyncio
import json
import sys
import time
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "backend"))
from app.env_fix import repair_ssl_env # noqa: E402
repair_ssl_env()
import httpx # noqa: E402
import ttc_core as core # noqa: E402
CACHE_DIR = REPO_ROOT / "run_cache"
def rebuild_messages(question: str, chain: list[dict]) -> list[dict]:
"""Reconstruct the conversation the original run held for this question.
Mirrors ttc_core.run_sequential_rounds exactly: system, question, the
baseline answer, then alternating REVISE_PROMPT and response for each
recorded round.
"""
messages = [
{"role": "system", "content": core.SYSTEM_PROMPT},
{"role": "user", "content": question},
{"role": "assistant", "content": chain[0]["text"]},
]
for entry in chain[1:]:
messages.append({"role": "user", "content": core.REVISE_PROMPT})
messages.append({"role": "assistant", "content": entry["text"]})
return messages
async def extend_one(client, sem, question, chain, gold, target, api_key, model):
messages = rebuild_messages(question, chain)
cumulative = chain[-1]["tokens"]
out = []
for r in range(len(chain) + 1, target + 1):
messages.append({"role": "user", "content": core.REVISE_PROMPT})
started = time.perf_counter()
text, tokens = await core.call_model(client, messages, 0.0, sem, api_key, model)
latency = time.perf_counter() - started
cumulative += tokens
messages.append({"role": "assistant", "content": text})
answer = core.extract_answer(text)
out.append({
"round": r,
"answer": answer,
"tokens": cumulative,
"call_tokens": tokens,
"correct": core.is_correct(answer, gold),
"latency_s": round(latency, 3),
})
return out
async def main() -> int:
args = [a for a in sys.argv[1:] if not a.startswith("--")]
target = next((int(a.split("=")[1]) for a in sys.argv[1:] if a.startswith("--rounds=")), 16)
if not args:
print(__doc__)
return 2
path = CACHE_DIR / f"{args[0]}.json"
if not path.exists():
print(f"No cached run at {path}", file=sys.stderr)
return 2
run = json.loads(path.read_text())
model = run["config"]["model"]
existing = run["config"]["sequential_rounds"]
api_key = [l.split("=", 1)[1].strip() for l in open(REPO_ROOT / "backend" / ".env")
if l.startswith("TOGETHER_API_KEY=")][0]
n = len(run["questions"])
new_rounds = target - existing
print(f"model : {model}")
print(f"questions : {n}")
print(f"rounds : K={existing} already recorded, extending to K={target}")
print(f"calls : {n * new_rounds}\n")
if new_rounds <= 0:
print("Nothing to extend.")
return 1
sem = asyncio.Semaphore(core.MAX_CONCURRENCY)
done = {"n": 0}
async with httpx.AsyncClient(timeout=120) as client:
async def one(i):
extra = await extend_one(
client, sem, run["questions"][i]["question"], run["sequential"][i],
run["questions"][i]["gold"], target, api_key, model)
done["n"] += 1
if done["n"] % 10 == 0:
print(f" {done['n']}/{n} chains extended", flush=True)
return extra
extended = await asyncio.gather(*[one(i) for i in range(n)])
out = REPO_ROOT / "fixtures" / f"sequential_ext_{args[0]}.json"
out.write_text(json.dumps({
"cache_key": args[0], "run_id": run["run_id"], "model": model,
"from_round": existing + 1, "to_round": target,
"chains": [
[{"round": e["round"], "answer": e["answer"], "tokens": e["tokens"],
"call_tokens": e["call_tokens"], "correct": e["correct"],
"latency_s": e["latency_s"]} for e in chain]
for chain in extended
],
}, indent=2))
print(f"\nwritten -> {out.relative_to(REPO_ROOT)}")
hard = [i for i, b in enumerate(run["baselines"]) if not b["correct"]]
for k in (4, 8, 12, 16):
acc = 0
for i in hard:
full = run["sequential"][i] + extended[i]
entry = next((e for e in full if e["round"] == k), None)
acc += bool(entry and entry["correct"])
print(f" hard subset accuracy at K={k:2}: {acc}/{len(hard)}")
return 0
if __name__ == "__main__":
raise SystemExit(asyncio.run(main()))