backend/app/ocr/quality.py

"""OCR output-quality + document-analysis benchmark.

For each available OCR backend, over a set of scanned samples with ground truth, we
measure two quality dimensions and capture logs + metrics:

  1. OCR text quality  — Character Error Rate (CER) and Word Error Rate (WER) of the
     transcribed text vs a reference (the `.txt` sidecar that ships with each scan).
  2. Document-analysis quality — field exact-match and F1 of the FULL pipeline
     (OCR → classify → extract → validate) vs the document's ground-truth JSON.

Plus latency, cost, model name/size. Results are published (file + endpoint + optional
HF upload). Pure-Python edit distance, no extra deps.
"""
from __future__ import annotations

import json
import re
import time
from pathlib import Path

from ..observability import log_event

# scanned samples that have BOTH a .txt sidecar (reference text) and a gt.json, AND
# that genuinely exercise each OCR engine independently (different CER per backend).
DEFAULT_SAMPLES = [
    "invoice_scanned_basic", "receipt_scanned", "po_scanned",
    "contract_scanned", "subscription_memo_scanned",
]
# field-accuracy-only (no sidecar reference text). The extreme tier
# (scripts/generate_extreme_docs.py — perspective photo, image collage, degraded fax) is a
# VISION-extraction stress set: on those images the OCR engines fall back to a shared text
# source (identical CER across backends), so they are excluded from the per-backend CER
# benchmark and scored on field accuracy only. complex_invoice_messy requires a real VLM.
FIELD_ONLY_SAMPLES = ["complex_invoice_messy"]


def _norm(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "").strip().lower())


def _lev(a, b) -> int:
    """Levenshtein distance over any sequences (str or list)."""
    if a == b:
        return 0
    la, lb = len(a), len(b)
    if la == 0:
        return lb
    if lb == 0:
        return la
    prev = list(range(lb + 1))
    for i in range(1, la + 1):
        cur = [i]
        ca = a[i - 1]
        for j in range(1, lb + 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (ca != b[j - 1])))
        prev = cur
    return prev[lb]


def cer(hyp: str, ref: str):
    ref, hyp = _norm(ref), _norm(hyp)
    if not ref:
        return None
    return round(min(1.0, _lev(hyp, ref) / len(ref)), 4)


def wer(hyp: str, ref: str):
    rw, hw = _norm(ref).split(), _norm(hyp).split()
    if not rw:
        return None
    return round(min(1.0, _lev(hw, rw) / len(rw)), 4)


def _model_size(settings, backend_name: str):
    """Map an OCR backend to its model name + size from the model registry."""
    from ..models_registry import model_catalog
    cat = model_catalog(settings)
    flat = [{**m, "lab": lab["lab"]} for lab in cat["labs"] for m in lab["models"]]
    if backend_name == "minicpm":
        m = next((x for x in flat if x["name"].startswith("MiniCPM-V")), None)
    elif backend_name == "cohere":
        m = next((x for x in flat if x["name"].startswith("Aya-Vision-8")), None)
    else:
        m = None
    if m:
        return {"model": m["name"], "params_b": m["params_b"], "size_gb": m["size_gb_int4"], "lab": m["lab"]}
    return {"model": backend_name, "params_b": None, "size_gb": None, "lab": "classic"}


def run_ocr_quality(settings, ocr_registry, router, metrics, db=None, rag_store=None,
                    samples=None) -> dict:
    from ..pipeline import process_document
    from evals import scorers

    samples = samples or DEFAULT_SAMPLES
    ds = settings.evals_dataset_dir
    available = [n for n in ocr_registry.available_names() if n != "sidecar"] + ["sidecar"]

    log_event("info", "OCR quality benchmark started",
              backends=available, samples=samples + FIELD_ONLY_SAMPLES)

    backend_rows = []
    for bname in available:
        backend = ocr_registry.get(bname)
        if not backend or not backend.available():
            continue
        cers, wers, exacts, f1s, lats, costs = [], [], [], [], [], []
        per_sample = []
        for sid in samples + FIELD_ONLY_SAMPLES:
            doc = _find(ds, sid)
            if not doc:
                continue
            gt = _load_gt(ds, sid)
            t0 = time.perf_counter()
            run = process_document(doc, router=router, settings=settings, metrics=metrics,
                                   ocr_registry=ocr_registry, ocr_backend=bname,
                                   db=db, rag_store=rag_store, doc_id=f"q-{bname}-{sid}",
                                   mode="quality")
            st = run["_state"]
            ocr_text = (st.get("ocr") or {}).get("text", "")
            ref = _ref_text(ds, sid)
            c = cer(ocr_text, ref) if ref else None
            w = wer(ocr_text, ref) if ref else None
            score = scorers.score_document(st.get("extracted") or {},
                                           {k: v for k, v in gt.items() if not k.startswith("_")}) if gt else {}
            case = {"sample": sid, "cer": c, "wer": w,
                    "field_exact": score.get("exact_match"), "field_f1": score.get("field_f1"),
                    "latency_ms": round((time.perf_counter() - t0) * 1000, 1),
                    "cost_usd": run.get("total_cost_usd", 0.0),
                    "confidence": st.get("confidence")}
            per_sample.append(case)
            if c is not None:
                cers.append(c); wers.append(w)
            if score.get("exact_match") is not None:
                exacts.append(score["exact_match"]); f1s.append(score["field_f1"] or 0)
            lats.append(case["latency_ms"]); costs.append(case["cost_usd"])
            log_event("info", f"OCR quality: {bname} on {sid}",
                      cer=c, wer=w, field_exact=score.get("exact_match"))

        size = _model_size(settings, bname)
        row = {
            "backend": bname, **size,
            "is_reference": bname == "sidecar",
            "cer": _avg(cers), "wer": _avg(wers),
            "field_exact_match": _avg(exacts), "field_f1": _avg(f1s),
            "avg_latency_ms": _avg(lats), "avg_cost_usd": round(_avg(costs) or 0, 6),
            "samples_scored": len(per_sample), "per_sample": per_sample,
        }
        backend_rows.append(row)

    # rank: best OCR text quality (lowest CER among real engines), best analysis (highest exact)
    real = [r for r in backend_rows if not r["is_reference"] and r["cer"] is not None]
    best_ocr = min(real, key=lambda r: r["cer"])["backend"] if real else None
    # rank only real engines for "best analysis" — sidecar is the reference text, not a
    # competing OCR backend, so it must never be crowned best.
    scored = [r for r in backend_rows
              if not r["is_reference"] and r["field_exact_match"] is not None]
    best_analysis = max(scored, key=lambda r: r["field_exact_match"])["backend"] if scored else None

    report = {
        "generated_at": time.time(),
        "note": "CER/WER vs .txt sidecar reference; field accuracy vs gt.json. "
                "sidecar = reference text source (CER≈0 by construction).",
        "models": _model_size_table(settings),
        "backends": backend_rows,
        "best_ocr_quality": best_ocr,
        "best_document_analysis": best_analysis,
    }
    if db is not None:
        try:
            db.audit("ocr_quality_benchmark",
                     detail={"best_ocr": best_ocr, "best_analysis": best_analysis,
                             "backends": [r["backend"] for r in backend_rows]})
        except Exception:
            pass
    log_event("info", "OCR quality benchmark complete",
              best_ocr=best_ocr, best_analysis=best_analysis)
    return report


def _model_size_table(settings):
    from ..models_registry import model_catalog
    return [{"lab": lab["lab"], **{k: m[k] for k in ("name", "params_b", "size_gb_int4", "modality", "role", "available")}}
            for lab in model_catalog(settings)["labs"] for m in lab["models"]]


def _avg(xs):
    xs = [x for x in xs if x is not None]
    return round(sum(xs) / len(xs), 4) if xs else None


def _find(ds: Path, sid: str):
    for ext in (".png", ".jpg", ".jpeg", ".pdf"):
        p = ds / f"{sid}{ext}"
        if p.exists():
            return p
    return None


def _ref_text(ds: Path, sid: str):
    p = ds / f"{sid}.txt"
    return p.read_text(encoding="utf-8", errors="ignore") if p.exists() else None


def _load_gt(ds: Path, sid: str):
    p = ds / f"{sid}.gt.json"
    return json.loads(p.read_text()) if p.exists() else None