lean-bench/leanbench/align.py

"""Compare the source-side simplex with the Lean-side one, role by role.

The result says *which* parameter disagrees -- a hypothesis the Lean dropped, one it
invented, a number system that changed, an ``∃`` turned into a ``∀`` -- which a
single "faithful: yes/no" from a judge cannot. Matching is deterministic and cheap:
shared symbols (``G``, ``n``) and shared concepts (``group``, ``finite``, ``prime``)
through a small synonym table. An LLM pairing (``llm_pairing``) can replace the
matching step, so the two can be compared as an ablation.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any

from .llm import ChatClient
from .project import Projection
from .schema import KIND_FACT, KIND_GOAL, KIND_OBJECT, KIND_OWN, Role
from .srcsimplex import SourceRole, SourceSimplex

DEFAULT_THRESHOLD = 0.3
GOAL_THRESHOLD = 0.2

ERROR_TYPES = (
    "missing_hypothesis", "missing_object", "extra_hypothesis", "extra_object",
    "type_conflict", "quantifier", "goal_mismatch", "degenerate",
)
# extra_object is reported but does not flag an item: Lean often introduces helpers.
FLAG_ERRORS = frozenset(ERROR_TYPES) - {"extra_object"}

CONCEPTS: dict[str, set[str]] = {
    "group": {"group", "grp"}, "abelian": {"abelian", "comm", "commutative"},
    "finite": {"finite", "fintype", "card"}, "infinite": {"infinite"},
    "ring": {"ring"}, "field": {"field"}, "domain": {"domain"},
    "integer": {"integer", "int", "ℤ"}, "natural": {"natural", "nat", "ℕ"},
    "real": {"real", "ℝ"}, "complex": {"complex", "ℂ"}, "rational": {"rational", "ℚ"},
    "prime": {"prime"}, "coprime": {"coprime", "relatively"},
    "divides": {"divides", "divisor", "dvd", "∣", "divisible", "divisibility"},
    "subgroup": {"subgroup"}, "ideal": {"ideal"}, "normal": {"normal"},
    "homomorphism": {"homomorphism", "hom", "morphism"},
    "isomorphism": {"isomorphism", "equiv", "≃", "≅", "isomorphic"},
    "set": {"set", "subset"}, "function": {"function", "map"},
    "even": {"even"}, "odd": {"odd"}, "cyclic": {"cyclic"}, "order": {"order", "orderof"},
    "gcd": {"gcd"}, "lcm": {"lcm"}, "positive": {"positive", "pos"},
    "nonzero": {"nonzero"}, "unit": {"unit", "units"}, "matrix": {"matrix"},
    "vector": {"vector"}, "space": {"space"}, "continuous": {"continuous"},
    "sequence": {"sequence", "seq"}, "polynomial": {"polynomial"},
    "kernel": {"kernel", "ker"}, "image": {"image", "range"},
    "injective": {"injective"}, "surjective": {"surjective"}, "bijective": {"bijective"},
}
_TOKEN_TO_CONCEPT = {tok: concept for concept, toks in CONCEPTS.items() for tok in toks}
NUMBER_SYSTEMS = {"integer", "natural", "real", "complex", "rational"}

_TOKEN = re.compile(r"[A-Z]+(?![a-z])|[A-Z]?[a-z]+|[ℕℤℚℝℂ∣≃≅]")
_IDENT = re.compile(r"[A-Za-z_][A-Za-z0-9_']*")


def tokens(text: str) -> set[str]:
    out = set()
    for raw in _TOKEN.findall(text):
        word = raw.lower()
        out.add(word)
        if word.endswith("s") and len(word) > 3:
            out.add(word[:-1])
    return out


def concepts(text_or_words: str | list[str]) -> set[str]:
    words = tokens(text_or_words) if isinstance(text_or_words, str) else {
        w for item in text_or_words for w in tokens(item)}
    return {_TOKEN_TO_CONCEPT[w] for w in words if w in _TOKEN_TO_CONCEPT}


@dataclass
class Unit:
    """One thing on the Lean side a source role can be matched to."""

    key: str
    kind: str               # object | fact | instance
    role: Role | None
    owner: str = ""
    symbols: set[str] = field(default_factory=set)
    concept_set: set[str] = field(default_factory=set)


@dataclass
class Match:
    source: str
    lean: str
    score: float
    how: str

    def as_dict(self) -> dict[str, Any]:
        return {"source": self.source, "lean": self.lean, "score": round(self.score, 3),
                "how": self.how}


@dataclass
class Conflict:
    kind: str
    detail: str
    source: str = ""
    lean: str = ""

    def as_dict(self) -> dict[str, str]:
        return {"kind": self.kind, "detail": self.detail, "source": self.source,
                "lean": self.lean}


@dataclass
class Alignment:
    matches: list[Match] = field(default_factory=list)
    conflicts: list[Conflict] = field(default_factory=list)
    score: float = 0.0
    flag_errors: frozenset[str] = FLAG_ERRORS

    @property
    def errors(self) -> list[str]:
        seen: list[str] = []
        for conflict in self.conflicts:
            if conflict.kind not in seen:
                seen.append(conflict.kind)
        return seen

    @property
    def flagged(self) -> bool:
        return any(c.kind in self.flag_errors for c in self.conflicts)

    def as_dict(self) -> dict[str, Any]:
        return {
            "score": round(self.score, 3), "flagged": self.flagged, "errors": self.errors,
            "matches": [m.as_dict() for m in self.matches],
            "conflicts": [c.as_dict() for c in self.conflicts],
        }


def lean_units(projection: Projection) -> list[Unit]:
    units: list[Unit] = []
    for role in projection.roles:
        if role.kind == KIND_OBJECT:
            classes = [str(c) for c in role.attrs.get("instances", [])]
            units.append(Unit(
                key=role.name, kind="object", role=role, symbols={role.name},
                concept_set=concepts(role.type_text) | concepts(" ".join(classes))))
            for text in classes:
                units.append(Unit(
                    key=f"inst:{role.name}:{text}", kind="instance", role=None,
                    owner=role.name, symbols={role.name}, concept_set=concepts(text)))
        elif role.kind == KIND_FACT:
            units.append(Unit(
                key=role.name, kind="fact", role=role,
                symbols=set(_IDENT.findall(role.type_text)) | {role.name},
                concept_set=concepts(role.type_text) | _relation_concepts(role.type_text)))
    return units


def _relation_concepts(text: str) -> set[str]:
    found: set[str] = set()
    if re.search(r"0\s*<|>\s*0", text):
        found.add("positive")
    if re.search(r"≠\s*0|0\s*≠", text):
        found.add("nonzero")
    if "∣" in text:
        found.add("divides")
    return found


def _score(source: SourceRole, unit: Unit) -> float:
    symbol = 1.0 if (set(source.symbols) & unit.symbols) else 0.0
    wanted = concepts(source.keywords) | concepts(source.description)
    if wanted and unit.concept_set:
        concept = len(wanted & unit.concept_set) / min(len(wanted), len(unit.concept_set))
    else:
        concept = 0.0
    return 0.6 * symbol + 0.4 * concept


def _compatible(source: SourceRole, unit: Unit) -> bool:
    if source.kind == "object":
        return unit.kind == "object"
    if source.kind == "fact":
        return unit.kind in ("fact", "instance")
    return False


def _pair_deterministic(
    source: SourceSimplex, units: list[Unit], threshold: float
) -> dict[str, tuple[Unit, float, str]]:
    candidates = []
    for role in source.roles:
        for unit in units:
            if _compatible(role, unit):
                score = _score(role, unit)
                if score >= threshold:
                    candidates.append((score, role.alias, unit.key, role, unit))
    candidates.sort(key=lambda c: (-c[0], c[1], c[2]))
    taken_sources: set[str] = set()
    taken_units: set[str] = set()
    pairs: dict[str, tuple[Unit, float, str]] = {}
    for score, alias, key, _role, unit in candidates:
        if alias in taken_sources:
            continue
        if key in taken_units and unit.kind != "instance":
            continue
        taken_sources.add(alias)
        taken_units.add(key)
        pairs[alias] = (unit, score, "symbol+concept")
    return pairs


def _existential_names(goal_text: str) -> set[str]:
    names: set[str] = set()
    for match in re.finditer(r"∃([^,]*),", goal_text):
        names.update(_IDENT.findall(match.group(1)))
    return names


def align(
    source: SourceSimplex,
    projection: Projection,
    *,
    threshold: float = DEFAULT_THRESHOLD,
    pairing: dict[str, str] | None = None,
    flag_errors: frozenset[str] = FLAG_ERRORS,
) -> Alignment:
    out = Alignment(flag_errors=flag_errors)
    if projection.degenerate:
        out.conflicts.append(Conflict("degenerate", "; ".join(projection.notes) or "no roles"))
    units = lean_units(projection)
    by_key = {u.key: u for u in units}

    if pairing is not None:
        pairs = {a: (by_key[k], 1.0, "llm") for a, k in pairing.items() if k in by_key}
    else:
        pairs = _pair_deterministic(source, units, threshold)

    goal_role = next((r for r in projection.roles if r.kind in (KIND_GOAL, KIND_OWN)), None)
    goal_text = goal_role.type_text if goal_role else ""
    exist_names = _existential_names(goal_text)

    used_units = {u.key for u, _s, _h in pairs.values()}
    matched_objects: dict[str, Unit] = {}
    for role in source.roles:
        if role.kind == "goal":
            continue
        hit = pairs.get(role.alias)
        if hit:
            unit, score, how = hit
            out.matches.append(Match(role.alias, unit.key, score, how))
            if role.kind == "object":
                matched_objects[role.alias] = unit
                _check_object(role, unit, exist_names, out)
            continue
        symbols = set(role.symbols)
        if role.kind == "object":
            if symbols & exist_names and role.introduced == "exists":
                out.matches.append(Match(role.alias, "∃-in-goal", 1.0, "existential"))
            elif symbols & exist_names:
                out.conflicts.append(Conflict(
                    "quantifier", f"{sorted(symbols & exist_names)} is ∃-bound in Lean "
                    f"but the source says '{role.introduced}'", role.alias))
            else:
                out.conflicts.append(Conflict("missing_object", role.description, role.alias))
        else:
            if _absorbed(role, units):
                out.matches.append(Match(role.alias, "type", 0.5, "absorbed"))
            else:
                out.conflicts.append(Conflict("missing_hypothesis", role.description, role.alias))

    for unit in units:
        if unit.kind == "fact" and unit.key not in used_units:
            out.conflicts.append(Conflict("extra_hypothesis", unit.role.type_text, "", unit.key))
        elif unit.kind == "object" and unit.key not in used_units:
            out.conflicts.append(Conflict(
                "extra_object", f"{unit.key} : {unit.role.type_text}", "", unit.key))

    src_goal = source.goal
    goal_ok = True
    if src_goal is not None and goal_role is not None:
        goal_symbols = set(_IDENT.findall(goal_text))
        sym = 1.0 if set(src_goal.symbols) & goal_symbols else 0.0
        wanted = concepts(src_goal.keywords) | concepts(src_goal.description)
        have = concepts(goal_text) | _relation_concepts(goal_text)
        con = len(wanted & have) / min(len(wanted), len(have)) if wanted and have else 0.0
        informative = bool(src_goal.symbols or wanted)
        if informative and 0.5 * sym + 0.5 * con < GOAL_THRESHOLD:
            goal_ok = False
            out.conflicts.append(Conflict(
                "goal_mismatch", f"source: {src_goal.description} | lean: {goal_text[:120]}",
                src_goal.alias, goal_role.name))
        lean_exists = goal_text.count("∃") > 0
        if source.goal_has_exists != lean_exists:
            out.conflicts.append(Conflict(
                "quantifier",
                f"source goal ∃={source.goal_has_exists}, lean goal ∃={lean_exists}",
                src_goal.alias, goal_role.name))
    if goal_ok and src_goal is not None and goal_role is not None:
        out.matches.append(Match(src_goal.alias, goal_role.name, 1.0, "goal"))

    n_source = len(source.roles)
    n_lean = len([u for u in units if u.kind != "instance"]) + (1 if goal_role else 0)
    out.score = 2 * len(out.matches) / (n_source + n_lean) if (n_source + n_lean) else 0.0
    return out


def _check_object(role: SourceRole, unit: Unit, exist_names: set[str], out: Alignment) -> None:
    wanted = (concepts(role.keywords) | concepts(role.description)) & NUMBER_SYSTEMS
    have = unit.concept_set & NUMBER_SYSTEMS
    if wanted and have and wanted.isdisjoint(have):
        out.conflicts.append(Conflict(
            "type_conflict", f"source {sorted(wanted)} vs lean {sorted(have)}",
            role.alias, unit.key))
    bound_exists = unit.key in exist_names
    if role.introduced == "exists" and not bound_exists:
        out.conflicts.append(Conflict(
            "quantifier", f"source asserts {unit.key} exists, lean takes it as given",
            role.alias, unit.key))


def _absorbed(role: SourceRole, units: list[Unit]) -> bool:
    wanted = concepts(role.keywords) | concepts(role.description)
    if not wanted:
        return False
    for unit in units:
        if unit.kind == "object" and wanted & unit.concept_set:
            if not role.symbols or set(role.symbols) & unit.symbols:
                return True
    return False


PAIRING_PROMPT = """\
You pair the logical parameters of a math passage with the binders of a formal
statement. For each SOURCE role, give the key of the LEAN item that expresses the
same parameter, or null if none does. Pair by meaning, not by spelling. A LEAN item can
be used once, except instance items. Reply with JSON: {"pairs": {"<source alias>":
"<lean key or null>"}}."""


def llm_pairing(
    client: ChatClient, source: SourceSimplex, projection: Projection
) -> dict[str, str]:
    """Ablation: let a model do the matching step (it sees roles, not the passage)."""
    units = lean_units(projection)
    listing = {
        "source": [r.as_dict() for r in source.roles if r.kind != "goal"],
        "lean": [
            {"key": u.key, "kind": u.kind,
             "text": (u.role.type_text if u.role else u.key.split(":", 2)[-1])}
            for u in units
        ],
    }
    import json
    reply = client.complete_json([  # type: ignore[attr-defined]
        {"role": "system", "content": PAIRING_PROMPT},
        {"role": "user", "content": json.dumps(listing, ensure_ascii=False)},
    ])
    pairs = reply.get("pairs", {}) if isinstance(reply, dict) else {}
    valid = {u.key for u in units}
    return {a: k for a, k in pairs.items() if isinstance(k, str) and k in valid}