src/gates.py

"""
Programmatic quality gates. Free, deterministic, zero judge variance.

Contract: a story that FAILS any hard gate gets total reward 0, regardless of
what the judge said. Gates run before the judge so we can skip paying for
obviously-broken generations.

This module exists because of a specific autopsy finding. The prior run's
completeness check was:

    enders = ".!?\\"'"
    text[-1] in enders

which uses straight quotes only. The policy emits curly U+201D constantly, so
properly-ended stories were scored 0.3 instead of 1.0. Worse, 100% of prior-run
eval samples were truncated mid-word by a max_new_tokens wall, and the judge
rubric capped quality at 4 for truncation -- collapsing within-group quality
variance to ~0, which in GRPO means zero advantage. The quality signal was
silently dead for the entire run. Everything here is built to make that
failure loud instead of silent.
"""

from __future__ import annotations

import math
import re
import unicodedata
from dataclasses import dataclass, asdict, field

# Sentence terminators, including the curly/CJK variants the prior run missed.
TERMINALS = set('.!?…' + '"”’\'' + '。！？' + ')]}*_')
# Characters allowed to trail a real terminator (markdown emphasis, quotes).
_TRAILING_DECOR = set('*_`"”’\'» )]}')

MIN_WORDS = 150
MAX_WORDS = 600
MAX_NGRAM_REPEAT_FRAC = 0.18   # frac of 4-grams that are repeats
MIN_CHAR_ENTROPY = 3.2         # bits/char; English prose sits ~4.0-4.4
MAX_NONASCII_FRAC = 0.08       # em-dashes/curly quotes are fine; CJK dumps are not
MAX_LINE_REPEAT_FRAC = 0.25


@dataclass
class GateResult:
    passed: bool
    completeness: float                    # 1.0 ended cleanly, 0.0 truncated
    n_words: int
    reasons: list[str] = field(default_factory=list)
    # diagnostics, logged not rewarded
    repeat_4gram_frac: float = 0.0
    char_entropy: float = 0.0
    nonascii_frac: float = 0.0
    hit_token_cap: bool = False

    def as_dict(self) -> dict:
        return asdict(self)


def _words(text: str) -> list[str]:
    return text.split()


def _norm_tokens(text: str) -> list[str]:
    return re.sub(r"[^\w\s]", " ", text.lower()).split()


def ends_cleanly(text: str) -> bool:
    """True if the story ends on a sentence terminator.

    Strips trailing markdown/quote decoration first, so `...over.*` and
    `...done."` and `...gone.”` all count. Also rejects the specific
    mid-word cutoff signature: last token is a bare word with no terminator.
    """
    t = text.rstrip()
    if not t:
        return False
    while t and t[-1] in _TRAILING_DECOR:
        t = t[:-1].rstrip()
        if t and t[-1] in '.!?…。！？':
            return True
    return bool(t) and t[-1] in '.!?…。！？'


def char_entropy(text: str) -> float:
    """Shannon entropy over characters, bits/char."""
    if not text:
        return 0.0
    counts: dict[str, int] = {}
    for ch in text:
        counts[ch] = counts.get(ch, 0) + 1
    n = len(text)
    return -sum((c / n) * math.log2(c / n) for c in counts.values())


def repeat_ngram_frac(text: str, n: int = 4) -> float:
    """Fraction of n-grams that are non-first occurrences. Catches loops."""
    toks = _norm_tokens(text)
    if len(toks) < n + 1:
        return 0.0
    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
    return 1.0 - (len(set(grams)) / len(grams))


def repeat_line_frac(text: str) -> float:
    """Fraction of non-trivial lines that are duplicates. Catches list loops."""
    lines = [l.strip() for l in text.splitlines() if len(l.strip()) > 12]
    if len(lines) < 4:
        return 0.0
    return 1.0 - (len(set(lines)) / len(lines))


def nonascii_frac(text: str) -> float:
    """Fraction of chars outside Latin/punctuation. Typographic marks exempt."""
    if not text:
        return 0.0
    bad = 0
    for ch in text:
        if ord(ch) < 128:
            continue
        if unicodedata.category(ch) in ("Pd", "Pi", "Pf", "Po", "Zs", "Sm"):
            continue  # em dash, curly quotes, ellipsis, nbsp
        bad += 1
    return bad / len(text)


def check(
    text: str,
    *,
    finish_reason: str | None = None,
    min_words: int = MIN_WORDS,
    max_words: int = MAX_WORDS,
) -> GateResult:
    """Run all gates. `finish_reason` from vLLM/OpenAI ('length' == hit cap).

    Passing `finish_reason` is strongly preferred: it detects the token-cap
    truncation that destroyed the prior run directly, rather than inferring it
    from punctuation.
    """
    text = (text or "").strip()
    reasons: list[str] = []

    n_words = len(_words(text))
    hit_cap = finish_reason == "length"
    complete = ends_cleanly(text) and not hit_cap

    rep4 = repeat_ngram_frac(text, 4)
    repline = repeat_line_frac(text)
    ent = char_entropy(text)
    na = nonascii_frac(text)

    if not text:
        reasons.append("empty")
    if hit_cap:
        reasons.append("hit_token_cap")
    if not ends_cleanly(text):
        reasons.append("no_terminal_punctuation")
    if n_words < min_words:
        reasons.append(f"too_short({n_words}<{min_words})")
    if n_words > max_words:
        reasons.append(f"too_long({n_words}>{max_words})")
    if rep4 > MAX_NGRAM_REPEAT_FRAC:
        reasons.append(f"ngram_loop({rep4:.2f})")
    if repline > MAX_LINE_REPEAT_FRAC:
        reasons.append(f"line_loop({repline:.2f})")
    if text and ent < MIN_CHAR_ENTROPY:
        reasons.append(f"low_entropy({ent:.2f})")
    if na > MAX_NONASCII_FRAC:
        reasons.append(f"nonascii({na:.2f})")

    return GateResult(
        passed=not reasons,
        completeness=1.0 if complete else 0.0,
        n_words=n_words,
        reasons=reasons,
        repeat_4gram_frac=rep4,
        char_entropy=ent,
        nonascii_frac=na,
        hit_token_cap=hit_cap,
    )