| """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 |
| 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") |
|
|
