src/diffusion_lm/quality.py

"""Programmatic generation quality: prose damage, and a reward over claim recaps.

The weights below encode a priority order rather than a tuned optimum: being right first,
not lying second, readable prose third, and brevity last and only relative to how much the
problem actually required. They are meant to be calibrated against judged samples before
driving any policy optimisation, since a reward this cheap is also cheap to game.
"""

from __future__ import annotations

import re

WORD = re.compile(r"[A-Za-z']+")

# Ordering, not tuning. Lying is charged PER ITEM rather than as a fraction of the record:
# a reader can work around clumsy prose or a missing line, but a fabricated plate or a
# superseded amount silently corrupts the report, so one of them has to outweigh several
# recalled fields. Verbosity is charged per fact carried, so a harder case may run longer.
WEIGHTS = {
    'recall': 1.0,
    'stale': -0.35,
    'invented': -0.35,
    'duplication': -0.2,
    'verbosity': -0.15,
}
# Thinking tokens per recalled fact treated as free before the verbosity term bites.
VERBOSITY_BUDGET = 40.0
# Lie counts saturate here so a wholly hallucinated report does not dominate the gradient.
MAX_LIES = 3.0


def prose_damage(text: str) -> dict[str, float]:
    """Adjacent duplicate rate and lexical diversity, the parallel-sampling signatures."""

    words = [word.lower() for word in WORD.findall(text)]
    pairs = max(1, len(words) - 1)
    duplicates = sum(a == b for a, b in zip(words, words[1:]))
    return {
        'duplicate_rate': duplicates / pairs,
        'lexical_diversity': len(set(words)) / max(1, len(words)),
        'words': len(words),
    }


def reward(report_score: dict[str, object], think_tokens: int, answer_text: str) -> dict:
    """Combine fact fidelity, prose damage and verbosity into one scalar plus its parts.

    ``report_score`` comes from :func:`diffusion_lm.claims.score_report`. Penalties are
    fractions of the fact count so a long claim and a short one stay comparable.
    """

    recalled = max(1, int(report_score['recalled']))
    damage = prose_damage(answer_text)
    per_fact = think_tokens / recalled
    parts = {
        'recall': float(report_score['recall']),
        'stale': min(MAX_LIES, len(report_score['stale_kept'])),
        'invented': min(MAX_LIES, len(report_score['invented'])),
        'duplication': min(1.0, damage['duplicate_rate'] * 20.0),
        'verbosity': max(0.0, per_fact - VERBOSITY_BUDGET) / VERBOSITY_BUDGET,
    }
    total = sum(WEIGHTS[name] * value for name, value in parts.items())
    return {
        'total': total,
        'parts': parts,
        'think_tokens_per_fact': per_fact,
        'duplicate_rate': damage['duplicate_rate'],
        'lexical_diversity': damage['lexical_diversity'],
    }