Space deploy: full app minus figures/ (not needed by the running Docker app; HF's binary-content check rejected those PNGs)
c140d3a Download tools/extend_sequential.py from ChandraPrakashBathula/test-time-compute: direct link, hf CLI and curl.
- Browser
- Download file 5.77 kB
-
https://huggingface.co/spaces/ChandraPrakashBathula/test-time-compute/resolve/main/tools/extend_sequential.py
- Command line
-
hf download hf://spaces/ChandraPrakashBathula/test-time-compute/tools/extend_sequential.py
-
curl -L -o extend_sequential.py https://huggingface.co/spaces/ChandraPrakashBathula/test-time-compute/resolve/main/tools/extend_sequential.py
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())) | |