File size: 4,188 Bytes
e1e980e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Shared pieces for fitting and answering the `decisions` config of JS/TS Code Review Decisions.
Every model script reads cases with `load_split` and writes answers with `PredictionWriter`, in the JSONL shape that
`score_decisions.py` scores: {"id", "answers": {name: System One answer}, "latency_ms"}.
"""
from __future__ import annotations
import json
import os
import random
import time
from pathlib import Path
import pyarrow.parquet as pq
HERE = Path(__file__).resolve().parent
DATA = Path(os.environ.get("DECISIONS_DATA", HERE.parent / "js-ts-code-review-decisions-3000"))
def load_split(split: str, limit: int | None = None, seed: int = 0) -> list[dict]:
"""Cases of one split: id, task_slug, state (task, language, code), questions, and gold label per question.
`limit` keeps whole tasks (all 4 variants) chosen at random, so a small run still has clean and flawed code."""
table = pq.read_table(DATA / "decisions" / f"{split}.parquet", columns=["id", "task_slug", "state", "questions", "gold"])
rows = table.to_pylist()
if limit is not None and limit < len(rows):
tasks = sorted({r["task_slug"] for r in rows})
random.Random(seed).shuffle(tasks)
keep = set(tasks[: max(1, limit // 4)])
rows = [r for r in rows if r["task_slug"] in keep]
return [{"id": r["id"], "task_slug": r["task_slug"], "state": json.loads(r["state"]),
"questions": json.loads(r["questions"]),
"gold": {name: g["label"] for name, g in json.loads(r["gold"]).items()}} for r in rows]
def options_of(question: dict) -> list[str]:
"""The option labels `score_decisions.py` scores, in order: true/false, choice keys, or level indexes."""
if question["type"] == "noul":
return ["true", "false"]
if question["type"] == "choice":
return list(question["criteria"])
return [str(i) for i in range(len(question["criteria"]))]
def gold_index(question: dict, label: str) -> int:
return options_of(question).index(label)
def to_answer(question: dict, probs: list[float]) -> dict:
"""A System One answer from probabilities in `options_of` order (a noul's are [true, false])."""
probs = [float(p) for p in probs]
if question["type"] == "noul":
return {"type": "noul", "noul": probs[0]}
return {"type": question["type"], "probabilities": dict(zip(options_of(question), probs))}
def state_text(state: dict) -> str:
"""One string for models that take plain text: the task, the language, then the code."""
return f"Task: {state['task']}\nLanguage: {state['language']}\n\n{state['code']}"
class PredictionWriter:
"""Appends one answered case per line and flushes it, so a stopped run resumes where it left off."""
def __init__(self, path: str | Path):
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.done: set[str] = set()
if self.path.exists():
for line in self.path.read_text().splitlines():
try:
self.done.add(json.loads(line)["id"])
except (json.JSONDecodeError, KeyError):
break # a line cut short by a stop; everything before it is kept
self._handle = self.path.open("a")
def pending(self, cases: list[dict]) -> list[dict]:
return [c for c in cases if c["id"] not in self.done]
def write(self, case_id: str, answers: dict, latency_ms: float | None = None) -> None:
record = {"id": case_id, "answers": answers}
if latency_ms is not None:
record["latency_ms"] = round(latency_ms, 3)
self._handle.write(json.dumps(record) + "\n")
self._handle.flush()
self.done.add(case_id)
def close(self) -> None:
self._handle.close()
def log(path: str | Path | None, **fields) -> None:
"""Prints a JSON line and appends it to `path` when given (training curves, timings)."""
line = json.dumps({"time": round(time.time(), 1), **fields})
print(line, flush=True)
if path:
Path(path).parent.mkdir(parents=True, exist_ok=True)
with open(path, "a") as handle:
handle.write(line + "\n")
|