backend/app/classification.py

"""Clause classification into CUAD-aligned categories.

Two backends behind one entry point (`classify_all`), selected by the
CLASSIFIER env var:

    rules     transparent keyword/regex classifier — fast, zero ML deps,
              easy to debug live
    zeroshot  cascade: the regexes below act as a high-recall *candidate*
              generator, and a local DeBERTa-v3 zero-shot NLI model
              (see zeroshot.py) confirms or rejects each (clause, category)
              candidate by meaning. Kills the classic keyword false
              positives ("audited financial statements" -> audit_rights)
              without paying full O(clauses x categories) inference.
    finetuned a DeBERTa-v3 classifier fine-tuned on CUAD (see finetuned.py /
              scripts/train_classifier.py) — direct multi-label prediction,
              highest accuracy when trained weights are present.
    legalbert/fusion  fine-tuned sentence classifier (legalbert_clf.py) over its
              12 CUAD categories, with zero-shot/rules filling the rest. Pick the
              shipped model via LEGALBERT_MODEL_DIR (DeBERTa 0.67 / LegalBERT 0.70).
              Falls back to the base classifier if no fine-tuned weights present.
    auto      (default) zeroshot when its deps + cached weights are
              present, else rules — mirrors CORE_LLM_BACKEND=auto.
"""
from __future__ import annotations

import os
import re

from .schema import Clause

# category key -> (CUAD label name, [patterns])
CATEGORY_PATTERNS: dict[str, tuple[str, list[str]]] = {
    "termination": ("Termination For Convenience", [
        r"\bterminat(e|ion|ed)\b", r"\bterminate for convenience\b",
    ]),
    "auto_renewal": ("Renewal Term", [
        r"\bautomatic(ally)? renew", r"\brenewal term\b", r"\bsuccessive (renewal )?(terms?|periods?)\b",
    ]),
    "term": ("Agreement Date / Expiration Date", [
        r"\binitial term\b", r"\bterm of this agreement\b", r"\bshall (commence|remain in (full )?force)\b",
    ]),
    "liability_cap": ("Cap On Liability", [
        r"\blimitation of liability\b", r"\bliab(le|ility)\b.*\b(cap|limit|exceed|aggregate)\b",
        r"\bin no event\b.*\bliab", r"\bunlimited liability\b",
    ]),
    "indemnification": ("Indemnification", [
        r"\bindemnif(y|ies|ication)\b", r"\bhold harmless\b", r"\bdefend\b.*\bclaims?\b",
    ]),
    "governing_law": ("Governing Law", [
        r"\bgoverning law\b", r"\bgoverned by (the )?laws?\b", r"\bjurisdiction\b",
    ]),
    "confidentiality": ("Confidentiality", [
        r"\bconfidential(ity)?\b", r"\bnon-?disclosure\b", r"\btrade secrets?\b",
    ]),
    "payment_terms": ("Payment Terms", [
        r"\bpayment\b", r"\binvoice[sd]?\b", r"\bnet\s+\d{1,3}\b", r"\bfees?\b.*\bpayable\b",
        r"\blate (payment|fee)\b", r"\binterest\b.*\b(per month|per annum|%)",
    ]),
    "sla": ("Service Levels", [
        r"\bservice level\b", r"\buptime\b", r"\bavailability\b.*\b\d{2}(\.\d+)?%",
        r"\bservice credits?\b", r"\bresponse time\b",
    ]),
    "data_protection": ("Data Protection / Privacy", [
        r"\bdata protection\b", r"\bpersonal (data|information)\b", r"\bGDPR\b",
        r"\bsecurity (breach|incident)\b", r"\bprocess(ing|or)\b.*\bdata\b",
    ]),
    "ip_ownership": ("IP Ownership Assignment", [
        r"\bintellectual property\b", r"\bownership\b.*\b(work product|deliverables)\b",
        r"\bwork[- ]for[- ]hire\b", r"\blicen[cs]e[sd]?\b",
    ]),
    "insurance": ("Insurance", [
        r"\binsurance\b", r"\bpolicy limits?\b", r"\bcommercial general liability\b",
    ]),
    "audit_rights": ("Audit Rights", [
        r"\baudit\b", r"\bbooks and records\b", r"\binspect(ion)?\b.*\brecords\b",
    ]),
    "assignment": ("Anti-Assignment", [
        r"\bassign(ment)?\b.*\b(consent|prior written)\b", r"\bnot (be )?assign(able|ed)?\b",
    ]),
    "exclusivity": ("Exclusivity", [
        # NOT a bare \bexclusive\b — "exclusive remedy", "exclusive jurisdiction"
        # and "exclusive of taxes" are everywhere and are not exclusivity clauses
        r"\bexclusivity\b", r"\bexclusively from\b",
        r"\bsole(ly)? (supplier|provider|source)\b",
        r"\bnon-?compete\b", r"\bshall not\b[^.]{0,80}\bengage any third party\b",
    ]),
    "force_majeure": ("Force Majeure", [
        r"\bforce majeure\b", r"\bbeyond (its|the) reasonable control\b",
    ]),
    "warranty": ("Warranty Duration", [
        r"\bwarrant(s|y|ies)\b", r"\bas is\b.*\bwithout warrant",
    ]),
    "notice": ("Notice Period", [
        r"\bnotices?\b.*\b(writing|written|address)\b", r"\bwritten notice\b",
    ]),
    "deliverables": ("Deliverables", [
        r"\bdeliverables?\b", r"\bstatement of work\b", r"\bmilestones?\b",
    ]),
}

# Extra candidate-only patterns for the zero-shot cascade. Too noisy to be
# rules-mode classifications on their own (bare "exclusive" matches "exclusive
# remedy" everywhere), but as *candidates* the NLI model rejects the noise, so
# here recall is all that matters.
BROAD_PATTERNS: dict[str, list[str]] = {
    # Bare \bexclusive\b is too noisy ("exclusive remedy", "exclusive
    # jurisdiction", "exclusive of taxes" appear in almost every contract).
    # The negative lookahead drops those known false-positive contexts so the
    # zeroshot confirm stage sees fewer spurious candidates.
    "exclusivity": [r"\bexclusive\b(?!\s*(?:of\b|remedy\b|remedies\b|jurisdiction\b))"],
    "assignment": [r"\bassign(s|ed|ment)?\b"],
    "termination": [r"\bexpir(e|es|ation)\b", r"\bcancel(lation)?\b"],
    "liability_cap": [r"\bconsequential\b", r"\bindirect\b.*\bdamages\b"],
    "auto_renewal": [r"\brenew(s|al|ed)?\b", r"\bextend(ed|s)?\b.*\bterm\b"],
}

_COMPILED = {
    key: (label, [re.compile(p, re.IGNORECASE) for p in pats])
    for key, (label, pats) in CATEGORY_PATTERNS.items()
}

_COMPILED_BROAD = {
    key: [re.compile(p, re.IGNORECASE) for p in pats]
    for key, pats in BROAD_PATTERNS.items()
}

MAX_CATEGORIES = 4
# chars of clause text around the first pattern hit fed to the NLI model as
# the premise (the model truncates at 512 tokens, so feeding the whole clause
# could push the matched language out of the window)
_CTX_BEFORE, _CTX_AFTER = 300, 900


def classify_clause(clause: Clause) -> list[str]:
    """Rules backend: return category keys. Heading matches count double."""
    text = clause.text
    heading = (clause.heading or "")
    cats: list[tuple[str, int]] = []
    for key, (_label, pats) in _COMPILED.items():
        score = 0
        for p in pats:
            if p.search(heading):
                score += 2
            if p.search(text):
                score += 1
        if score >= 1:
            cats.append((key, score))
    cats.sort(key=lambda kv: -kv[1])
    return [k for k, _ in cats[:MAX_CATEGORIES]]


def _candidates(clause: Clause) -> list[tuple[str, str]]:
    """(category_key, premise_text) candidates for the zero-shot confirm
    stage — every category whose strict or broad patterns hit, uncapped."""
    text = clause.text
    heading = (clause.heading or "")
    out: list[tuple[str, str]] = []
    for key, (_label, pats) in _COMPILED.items():
        all_pats = pats + _COMPILED_BROAD.get(key, [])
        first = None
        for p in all_pats:
            m = p.search(text)
            if m and (first is None or m.start() < first):
                first = m.start()
        if first is None and not any(p.search(heading) for p in all_pats):
            continue
        start = max(0, (first or 0) - _CTX_BEFORE)
        premise = (heading + "\n" if heading else "") + text[start:(first or 0) + _CTX_AFTER]
        out.append((key, premise))
    return out


def classify_all(clauses: list[Clause]) -> str:
    """Classify every clause in place; returns the backend name used."""
    mode = os.environ.get("CLASSIFIER", "auto").lower()
    if mode == "finetuned":
        from . import finetuned
        fclf = finetuned.get_classifier()
        if fclf is None:
            raise RuntimeError(
                "CLASSIFIER=finetuned but no trained model was found. Train one "
                "with scripts/train_classifier.py (or set FINETUNED_MODEL_DIR), "
                "or use CLASSIFIER=zeroshot / rules.")
        fclf.classify(clauses)
        return "finetuned"
    if mode in ("legalbert", "fusion"):
        # Fine-tuned sentence classifier (LegalBERT 0.70 / DeBERTa 0.67 via
        # LEGALBERT_MODEL_DIR). It owns its 12 CUAD categories; the base
        # classifier (zero-shot DeBERTa, else rules) fills the rest, so every
        # category is covered and the dynamic baseline runs downstream unchanged.
        # Falls back gracefully to the base alone if the fine-tuned weights
        # aren't present — so the pipeline always works.
        from . import legalbert_clf
        lb = legalbert_clf.get_classifier()
        base = _zeroshot_or_rules(clauses, "auto")
        if lb is None:
            return base  # no fine-tuned weights yet -> base classifier only
        for c in clauses:
            ft = legalbert_clf.predict(lb, c)
            kept = [k for k in c.categories if k not in legalbert_clf.MODEL_CATS_SET]
            c.categories = (sorted(ft) + kept)[:MAX_CATEGORIES]
        return f"fusion({base}+finetuned)"
    return _zeroshot_or_rules(clauses, mode)


def _zeroshot_or_rules(clauses: list[Clause], mode: str) -> str:
    """Zero-shot DeBERTa cascade if available, else the regex classifier."""
    clf = None
    if mode in ("zeroshot", "auto"):
        from . import zeroshot
        clf = zeroshot.get_classifier(allow_download=(mode == "zeroshot"))
        if clf is None and mode == "zeroshot":
            raise RuntimeError(
                "CLASSIFIER=zeroshot but the model could not be loaded — "
                "install backend/requirements-ml.txt (torch, transformers, "
                "sentencepiece, protobuf) or set CLASSIFIER=rules.")
    if clf is None:
        for c in clauses:
            c.categories = classify_clause(c)
        return "rules"

    pairs: list[tuple[str, str]] = []   # (premise, key) for the model
    index: list[tuple[int, str]] = []   # (clause_idx, key) aligned with pairs
    for ci, c in enumerate(clauses):
        for key, premise in _candidates(c):
            pairs.append((premise, key))
            index.append((ci, key))

    probs = clf.entail_probs(pairs)
    from .zeroshot import THRESHOLD
    confirmed: dict[int, list[tuple[str, float]]] = {}
    for (ci, key), p in zip(index, probs):
        if p >= THRESHOLD:
            confirmed.setdefault(ci, []).append((key, p))
    for ci, c in enumerate(clauses):
        cats = sorted(confirmed.get(ci, []), key=lambda kp: -kp[1])
        c.categories = [k for k, _ in cats[:MAX_CATEGORIES]]
    return "zeroshot"


def cuad_label(category_key: str) -> str:
    return _COMPILED[category_key][0] if category_key in _COMPILED else category_key