""" Pluggable LLM judge with aggressive on-disk caching. Design decisions worth stating, because they differ from the prior run: 1. ONE STORY PER CALL, absolute rubric. The prior run sent all 16 samples in a single call and asked for relative scores. That is cheaper but wrong here: - quality is compared against a fixed gate tau (5.0), so it must be absolutely calibrated, not calibrated relative to whatever else happened to be in the group; - eval compares quality ACROSS models, which relative scoring makes meaningless; - batched scoring has position bias. Per-story also makes the cache actually work: identical stories recur across arms and eval passes, and each is paid for once, ever. 2. NO group_diversity FIELD. A set-level scalar is constant within a GRPO group -> zero within-group std -> zero advantage. It cannot train anything. All diversity credit comes from embeddings (src/diversity.py), where it is per-sample by construction. 3. Judge novelty is an ABSOLUTE 'freshness vs. generic AI slop' axis, not 'different from the others'. The latter is the embeddings' job. This keeps the two signals from being redundant-but-inconsistent. """ from __future__ import annotations import asyncio import hashlib import json import os import sqlite3 import time from dataclasses import dataclass, asdict from pathlib import Path RUBRIC_VERSION = "v4" SYSTEM_RUBRIC = """You are a demanding fiction editor scoring ONE short story written to a prompt. Score on absolute standards, not relative to anything you have seen before. STEP 1 -- FINISHEDNESS. Read the LAST sentence first, before anything else. A story is UNFINISHED if the final sentence stops mid-clause or mid-word, or if the piece simply halts with the situation unresolved and no closing beat. Beautiful prose does not make an unfinished story finished. If UNFINISHED: set "ended":false, quality 1-3, novelty 0-3, and stop deliberating -- an unfinished fragment is never novel, because you cannot know what it became. Only if it is FINISHED, continue to step 2. STEP 2 -- quality (0-10): coherence + prompt fulfilment + craft 0-2 Incoherent, or ignores the prompt entirely. 3-4 Followable but generic: stock premise, flat prose, perfunctory ending. 5-6 Competent. Does what the prompt asks, reads cleanly, resolves. Forgettable. 7-8 Strong. Controlled voice, purposeful structure, an ending that lands. 9-10 Excellent. Precise images, earned emotional turn, nothing wasted. If it never engages the prompt's actual premise, quality <= 4. STEP 3 -- novelty (0-10): freshness of premise, voice and construction vs. generic AI fiction 0-2 Pure slop signature: portentous one-line paragraphs, "It wasn't X, it was Y", italicised abstract nouns doing the emotional work, cosmic-melancholy fog. 3-4 Familiar premise handled in the familiar way. 5-6 One genuine choice (an unusual angle OR a distinct voice OR an odd structure). 7-8 Several: takes a real swing on premise, form or register and controls it. 9-10 Startling and still coherent. You could not have predicted this. Novelty is NOT weirdness: incoherence, non-sequiturs and word salad score 0-2. Judge freshness against fiction in general, NOT against other samples. Return ONLY minified JSON, no markdown fence: {"quality":<0-10>,"novelty":<0-10>,"ended":,"note":"<=10 words"}""" @dataclass class JudgeScore: quality: float novelty: float ended: bool = True note: str = "" ok: bool = True # False => call failed, scores are neutral fallbacks cached: bool = False def as_dict(self) -> dict: return asdict(self) NEUTRAL = JudgeScore(quality=5.0, novelty=5.0, ended=True, note="judge_failed", ok=False) # --------------------------------------------------------------- disk cache class ScoreCache: """sqlite keyed by sha256(model|rubric|prompt|story). Survives restarts.""" def __init__(self, path: str | Path): self.path = str(path) Path(self.path).parent.mkdir(parents=True, exist_ok=True) self._db = sqlite3.connect(self.path, check_same_thread=False) self._db.execute("PRAGMA journal_mode=WAL") self._db.execute( "CREATE TABLE IF NOT EXISTS scores (" " k TEXT PRIMARY KEY, quality REAL, novelty REAL," " ended INT, note TEXT, ts REAL)" ) self._db.commit() self.hits = 0 self.misses = 0 @staticmethod def key(model: str, prompt: str, story: str) -> str: h = hashlib.sha256() h.update(f"{model}\x00{RUBRIC_VERSION}\x00{prompt}\x00{story}".encode()) return h.hexdigest() def get(self, k: str) -> JudgeScore | None: r = self._db.execute( "SELECT quality,novelty,ended,note FROM scores WHERE k=?", (k,) ).fetchone() if r is None: self.misses += 1 return None self.hits += 1 return JudgeScore(quality=r[0], novelty=r[1], ended=bool(r[2]), note=r[3] or "", ok=True, cached=True) def put(self, k: str, s: JudgeScore) -> None: if not s.ok: return # never cache a failure self._db.execute( "INSERT OR REPLACE INTO scores VALUES (?,?,?,?,?,?)", (k, s.quality, s.novelty, int(s.ended), s.note[:120], time.time()), ) self._db.commit() def stats(self) -> dict: n = self._db.execute("SELECT COUNT(*) FROM scores").fetchone()[0] tot = self.hits + self.misses return {"rows": n, "hits": self.hits, "misses": self.misses, "hit_rate": (self.hits / tot) if tot else 0.0} # ------------------------------------------------------------------- judges def _parse(text: str) -> JudgeScore: t = (text or "").strip() if t.startswith("```"): t = t.split("\n", 1)[-1].rsplit("```", 1)[0].strip() i, j = t.find("{"), t.rfind("}") if i >= 0 and j > i: t = t[i:j + 1] d = json.loads(t) def num(v, lo=0.0, hi=10.0): return max(lo, min(hi, float(v))) return JudgeScore( quality=num(d.get("quality", 5)), novelty=num(d.get("novelty", d.get("novel", 5))), ended=bool(d.get("ended", True)), note=str(d.get("note", ""))[:120], ok=True, ) class OpenRouterJudge: """Async judge over OpenRouter. Default backend.""" BASE = "https://openrouter.ai/api/v1/chat/completions" def __init__( self, model: str = "deepseek/deepseek-v4-flash-0731", api_key: str | None = None, cache_path: str | Path = "cache/judge.sqlite", concurrency: int = 12, max_retries: int = 4, temperature: float = 0.0, timeout: float = 120.0, ): self.model = model self.api_key = api_key or os.environ.get("OPENROUTER_API_KEY", "") if not self.api_key: raise RuntimeError("OPENROUTER_API_KEY not set") self.cache = ScoreCache(cache_path) # NOT created here. asyncio.Semaphore binds to the loop that first # awaits it, and score_many_sync() calls asyncio.run() -> a NEW loop # every training step. A semaphore built in __init__ survives step 1 and # then raises "bound to a different event loop" on step 2, which our # retry wrapper would convert into neutral 5.0 scores forever. Built # per-call in score_many() instead. self.concurrency = concurrency self.sem: asyncio.Semaphore | None = None self.max_retries = max_retries self.temperature = temperature self.timeout = timeout self.n_calls = 0 self.tok_in = 0 self.tok_out = 0 self.n_failed = 0 self.n_empty = 0 self.n_truncated = 0 self.last_error = "" def _user(self, prompt: str, story: str) -> str: return f"WRITING PROMPT:\n{prompt}\n\nSTORY:\n{story}\n\nJSON only:" async def _one(self, client, prompt: str, story: str) -> JudgeScore: payload = { "model": self.model, "messages": [ {"role": "system", "content": SYSTEM_RUBRIC}, {"role": "user", "content": self._user(prompt, story)}, ], "temperature": self.temperature, # deepseek-v4-flash-0731 is a REASONING model. Left alone it spends # its whole output budget on `reasoning` and returns content=null # with finish_reason="length" -- which our fallback silently turned # into a neutral 5.0, i.e. a judge that scored everything identically # while looking like it worked. Calibration caught it at 4/166 calls # succeeding. Disabling reasoning is also 25x cheaper and 5x faster # (30 output tokens vs 764). "reasoning": {"enabled": False}, # Generous bound on a structured output that empirically needs ~30 # tokens. Not a length cap on content we care about -- and if it is # ever hit, `n_truncated` below makes it loud instead of silent. "max_tokens": 300, } delay = 2.0 for attempt in range(self.max_retries): try: async with self.sem: r = await client.post( self.BASE, json=payload, timeout=self.timeout, headers={"Authorization": f"Bearer {self.api_key}"}, ) if r.status_code in (429, 500, 502, 503, 529): await asyncio.sleep(delay); delay *= 2 continue r.raise_for_status() body = r.json() self.n_calls += 1 u = body.get("usage") or {} self.tok_in += int(u.get("prompt_tokens", 0) or 0) self.tok_out += int(u.get("completion_tokens", 0) or 0) ch = body["choices"][0] if ch.get("finish_reason") == "length": self.n_truncated += 1 content = ch["message"].get("content") if not content: self.n_empty += 1 raise ValueError("empty content (reasoning ate the budget?)") return _parse(content) except Exception as e: self.last_error = f"{type(e).__name__}: {str(e)[:140]}" if attempt == self.max_retries - 1: self.n_failed += 1 return NEUTRAL await asyncio.sleep(delay); delay *= 2 self.n_failed += 1 return NEUTRAL def health(self) -> dict: """Fail-loud accounting. A judge that returns neutral 5.0 for every call looks exactly like a judge that works, until you check this.""" att = self.n_calls + self.n_failed return {"calls_ok": self.n_calls, "failed": self.n_failed, "empty_content": self.n_empty, "truncated": self.n_truncated, "fail_rate": (self.n_failed / att) if att else 0.0, "last_error": self.last_error} def assert_healthy(self, max_fail_rate: float = 0.05) -> None: h = self.health() if h["fail_rate"] > max_fail_rate: raise RuntimeError( f"JUDGE UNHEALTHY: {h}. Refusing to train on neutral fallbacks -- " f"a constant reward column produces zero GRPO advantage." ) async def score_many(self, pairs: list[tuple[str, str]]) -> list[JudgeScore]: """pairs = [(prompt, story), ...] -> scores in the same order.""" import httpx out: list[JudgeScore | None] = [None] * len(pairs) todo: list[int] = [] for i, (p, s) in enumerate(pairs): hit = self.cache.get(self.cache.key(self.model, p, s)) if hit is not None: out[i] = hit else: todo.append(i) if todo: self.sem = asyncio.Semaphore(self.concurrency) # bind to THIS loop async def one_and_cache(i): """Persist each score the moment it lands, not after the whole batch. The pool stage issues ~16k calls in a single score_many(); caching only at the end meant a crash at 99% threw away every call. Incremental writes make any failure resumable for free -- a re-run replays as cache hits.""" s = await self._one(client, pairs[i][0], pairs[i][1]) if s.ok: self.cache.put(self.cache.key(self.model, *pairs[i]), s) return i, s async with httpx.AsyncClient(http2=False) as client: for fut in asyncio.as_completed([one_and_cache(i) for i in todo]): i, s = await fut out[i] = s return [o if o is not None else NEUTRAL for o in out] def score_many_sync(self, pairs: list[tuple[str, str]]) -> list[JudgeScore]: """Blocking wrapper for use inside TRL reward functions.""" try: asyncio.get_running_loop() except RuntimeError: return asyncio.run(self.score_many(pairs)) # already inside a loop (rare in TRL): run in a private one import concurrent.futures with concurrent.futures.ThreadPoolExecutor(1) as ex: return ex.submit(asyncio.run, self.score_many(pairs)).result() def cost_estimate(self, price_in: float, price_out: float) -> dict: """price_* in $ per 1M tokens.""" c = self.tok_in / 1e6 * price_in + self.tok_out / 1e6 * price_out return {"calls": self.n_calls, "tok_in": self.tok_in, "tok_out": self.tok_out, "usd": round(c, 4), **self.cache.stats()} def build_judge(**kw) -> OpenRouterJudge: """Factory. Extend here if an Anthropic/OpenAI key shows up.""" return OpenRouterJudge(**kw)