quantization/layer_sensitivity.py

"""Layer sensitivity analysis for OCR-aware mixed-precision quantization.

This script quantizes one layer group at a time, runs inference on a small
evaluation set, and measures the quality delta (CER, digit error, table structure)
relative to the BF16 baseline.

Usage:
    python quantization/layer_sensitivity.py \
        --model-path baidu/Unlimited-OCR \
        --eval-dir ./eval_images/ \
        --output sensitivity_results.json

Requires: mlx-vlm, mlx, Pillow, numpy
"""

from __future__ import annotations

import argparse
from collections import Counter
import gc
import json
import sys
import time
from pathlib import Path

import numpy as np

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from benchmarks.evaluate_cer import compute_cer as _compute_cer  # noqa: E402
from benchmarks.evaluate_cer import compute_digit_cer as _compute_digit_cer  # noqa: E402
from benchmarks.evaluate_tables import evaluate_tables  # noqa: E402
from benchmarks.normalize_output import normalize_ocr_output  # noqa: E402
from quantization.mixed_precision_convert import matches_pattern  # noqa: E402
from quantization.release_gate import dataset_manifest  # noqa: E402
from unlimited_ocr.engine import SlidingWindowNoRepeatNGramProcessor  # noqa: E402


# Quantizable MLX module groups to test independently. Normalization and the
# MoE router are raw arrays rather than quantizable MLX layers, so they remain
# BF16 by construction and are not presented as runnable sensitivity groups.
LAYER_GROUPS: dict[str, tuple[str, ...]] = {
    "vision_encoder": ("vision_model", "sam_model"),
    "vision_projector": ("projector",),
    "token_embeddings": ("language_model.model.embed_tokens",),
    "attention_q_proj": ("language_model.model.layers.*.self_attn.q_proj",),
    "attention_k_proj": ("language_model.model.layers.*.self_attn.k_proj",),
    "attention_v_proj": ("language_model.model.layers.*.self_attn.v_proj",),
    "attention_o_proj": ("language_model.model.layers.*.self_attn.o_proj",),
    "dense_mlp": (
        "language_model.model.layers.*.mlp.gate_proj",
        "language_model.model.layers.*.mlp.up_proj",
        "language_model.model.layers.*.mlp.down_proj",
    ),
    "shared_experts": ("language_model.model.layers.*.mlp.shared_experts",),
    "routed_experts": ("language_model.model.layers.*.mlp.switch_mlp",),
    "lm_head": ("language_model.lm_head",),
}

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".tiff", ".tif", ".webp", ".bmp"}


def collect_evaluation_files(
    eval_dir: Path,
    prompt: str = "<image>document parsing.",
    max_tokens: int = 4096,
) -> list[tuple[Path, Path]]:
    """Validate an evaluation dataset and return complete image/text pairs."""
    images_dir = eval_dir / "images"
    gt_dir = eval_dir / "ground_truth"
    if not images_dir.is_dir():
        raise FileNotFoundError(f"Images directory not found: {images_dir}")
    if not gt_dir.is_dir():
        raise FileNotFoundError(f"Ground-truth directory not found: {gt_dir}")
    if not isinstance(prompt, str) or prompt.count("<image>") != 1:
        raise ValueError("prompt must contain exactly one literal '<image>' token")
    if not isinstance(max_tokens, int) or isinstance(max_tokens, bool) or max_tokens < 1:
        raise ValueError("max_tokens must be a positive integer")

    image_files = sorted(
        path for path in images_dir.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
    )
    if not image_files:
        raise ValueError(f"No supported images found in {images_dir}")

    duplicate_stems = sorted(
        stem for stem, count in Counter(path.stem for path in image_files).items()
        if count > 1
    )
    if duplicate_stems:
        raise ValueError(
            "Multiple input images share the same stem: "
            + ", ".join(duplicate_stems)
        )

    missing = [
        path.name for path in image_files
        if not (gt_dir / f"{path.stem}.txt").is_file()
    ]
    if missing:
        raise FileNotFoundError(
            "Missing ground-truth text for: " + ", ".join(missing)
        )
    return [(path, gt_dir / f"{path.stem}.txt") for path in image_files]


def compute_cer(reference: str, hypothesis: str) -> float:
    """Compute Character Error Rate with the shared memory-efficient metric."""
    return _compute_cer(reference, hypothesis)


def compute_digit_cer(reference: str, hypothesis: str) -> float:
    """Compute CER only on digit characters."""
    return _compute_digit_cer(reference, hypothesis)


def run_inference(model, processor, image_path: str, prompt: str, max_tokens: int = 4096) -> str:
    """Run a single inference and return the text output."""
    from mlx_vlm import generate

    response = generate(
        model, processor,
        prompt=prompt,
        image=[image_path],
        max_tokens=max_tokens,
        temperature=0.0,
        logits_processors=[SlidingWindowNoRepeatNGramProcessor(35, 128)],
        verbose=False,
    )

    if isinstance(response, str):
        return response
    elif hasattr(response, "text"):
        return response.text
    return str(response)


def evaluate_model(
    model,
    processor,
    eval_dir: Path,
    prompt: str = "<image>document parsing.",
    max_tokens: int = 4096,
) -> dict:
    """Evaluate model on a directory of images with ground truth.

    Expects eval_dir to contain:
      - images/ subdirectory with input images
      - ground_truth/ subdirectory with .txt files (same stem as image)
    """
    evaluation_files = collect_evaluation_files(eval_dir, prompt, max_tokens)

    cer_scores = []
    digit_cer_scores = []
    table_scores = []
    total_time = 0.0

    for img_file, gt_file in evaluation_files:
        reference = normalize_ocr_output(gt_file.read_text(encoding="utf-8"))

        start = time.perf_counter()
        hypothesis = normalize_ocr_output(run_inference(
            model,
            processor,
            str(img_file),
            prompt,
            max_tokens=max_tokens,
        ))
        elapsed = time.perf_counter() - start

        cer = compute_cer(reference, hypothesis)
        digit_cer = compute_digit_cer(reference, hypothesis)

        cer_scores.append(cer)
        if any(char.isascii() and char.isdigit() for char in reference):
            digit_cer_scores.append(digit_cer)
        table_result = evaluate_tables(reference, hypothesis)
        if table_result["num_ref_tables"] > 0:
            table_scores.append(table_result["mean_score"])
        total_time += elapsed

    return {
        "num_samples": len(cer_scores),
        "num_digit_samples": len(digit_cer_scores),
        "num_table_samples": len(table_scores),
        "mean_cer": float(np.mean(cer_scores)) if cer_scores else 0.0,
        "mean_digit_cer": float(np.mean(digit_cer_scores)) if digit_cer_scores else None,
        "mean_table_score": float(np.mean(table_scores)) if table_scores else None,
        "total_inference_time": total_time,
        "avg_time_per_image": total_time / max(len(cer_scores), 1),
    }


def quantize_layer_group(model, config: dict, group_name: str) -> list[str]:
    """Quantize one configured MLX module group to MXFP8 in place.

    Returns the concrete module paths that were converted.
    """
    if group_name not in LAYER_GROUPS:
        available = ", ".join(LAYER_GROUPS)
        raise ValueError(f"Unknown layer group '{group_name}'. Available: {available}")

    from mlx_vlm.quant_utils import quantize_model

    patterns = LAYER_GROUPS[group_name]
    matched: list[str] = []

    def predicate(path, module):
        if any(matches_pattern(path, pattern) for pattern in patterns):
            matched.append(path)
            return {"group_size": 32, "bits": 8, "mode": "mxfp8"}
        return False

    quantize_model(
        model,
        config,
        group_size=32,
        bits=8,
        mode="mxfp8",
        quant_predicate=predicate,
    )
    if not matched:
        raise RuntimeError(
            f"Layer group '{group_name}' matched no quantizable MLX modules"
        )
    return sorted(set(matched))


def main():
    parser = argparse.ArgumentParser(
        description="Layer sensitivity analysis for OCR-aware quantization"
    )
    parser.add_argument("--model-path", required=True,
                        help="Path to BF16 model (baidu/Unlimited-OCR or local)")
    parser.add_argument(
        "--source-id",
        default=None,
        help="Public source identifier recorded in results (defaults to repo ID or local name)",
    )
    parser.add_argument(
        "--source-revision",
        default=None,
        help="Immutable source commit used for remote loading and provenance",
    )
    parser.add_argument("--eval-dir", required=True, type=Path,
                        help="Evaluation directory with images/ and ground_truth/")
    parser.add_argument("--output", type=Path, default=Path("sensitivity_results.json"),
                        help="Output JSON file for results")
    parser.add_argument("--prompt", default="<image>document parsing.",
                        help="OCR prompt to use")
    parser.add_argument("--max-tokens", type=int, default=4096)
    parser.add_argument(
        "--groups",
        nargs="+",
        choices=list(LAYER_GROUPS),
        default=list(LAYER_GROUPS),
        help="Layer groups to test (default: all)",
    )
    args = parser.parse_args()
    if args.max_tokens < 1:
        parser.error("--max-tokens must be positive")

    # Fail before importing MLX or allocating the model for dataset mistakes.
    collect_evaluation_files(args.eval_dir, args.prompt, args.max_tokens)
    dataset = dataset_manifest(args.eval_dir)

    print("=" * 60)
    print("Layer Sensitivity Analysis for Unlimited-OCR MLX")
    print("=" * 60)

    # Step 1: Load BF16 baseline
    print("\n[1/3] Loading BF16 baseline model...")
    from mlx_vlm import load
    from mlx_vlm.utils import load_config
    import mlx.core as mx

    base_config = load_config(args.model_path, revision=args.source_revision)
    if base_config.get("quantization") or base_config.get("quantization_config"):
        raise ValueError(
            "Layer sensitivity requires an unquantized BF16 source model"
        )
    model, processor = load(args.model_path, revision=args.source_revision)

    # Step 2: Evaluate baseline
    print("\n[2/3] Evaluating BF16 baseline...")
    baseline = evaluate_model(
        model,
        processor,
        args.eval_dir,
        args.prompt,
        max_tokens=args.max_tokens,
    )
    print(f"  Baseline CER: {baseline['mean_cer']:.4f}")
    digit_label = (
        f"{baseline['mean_digit_cer']:.4f}"
        if baseline["mean_digit_cer"] is not None
        else "N/A"
    )
    print(f"  Baseline Digit CER: {digit_label}")
    print(f"  Samples: {baseline['num_samples']}")
    model = None
    processor = None
    gc.collect()
    mx.clear_cache()

    # Step 3: Per-group sensitivity
    print("\n[3/3] Testing per-group quantization sensitivity...")
    results = {
        "baseline": baseline,
        "model_path": args.source_id or (
            Path(args.model_path).name if Path(args.model_path).is_dir() else args.model_path
        ),
        "source_revision": args.source_revision,
        "eval_dir": args.eval_dir.name,
        "dataset": dataset,
        "prompt": args.prompt,
        "max_tokens": args.max_tokens,
        "groups_tested": args.groups,
        "layer_groups": {},
    }

    for group in args.groups:
        print(f"\n  Testing: {group}")
        group_model = None
        group_processor = None
        try:
            group_model, group_processor = load(
                args.model_path,
                revision=args.source_revision,
            )
            matched_modules = quantize_layer_group(
                group_model,
                dict(base_config),
                group,
            )
            metrics = evaluate_model(
                group_model,
                group_processor,
                args.eval_dir,
                args.prompt,
                max_tokens=args.max_tokens,
            )
            digit_delta = (
                metrics["mean_digit_cer"] - baseline["mean_digit_cer"]
                if metrics["mean_digit_cer"] is not None
                and baseline["mean_digit_cer"] is not None
                else None
            )
            table_delta = (
                baseline["mean_table_score"] - metrics["mean_table_score"]
                if metrics["mean_table_score"] is not None
                and baseline["mean_table_score"] is not None
                else None
            )
            results["layer_groups"][group] = {
                "status": "success",
                "matched_module_count": len(matched_modules),
                "matched_modules": matched_modules,
                "metrics": metrics,
                "cer_delta": metrics["mean_cer"] - baseline["mean_cer"],
                "digit_cer_delta": digit_delta,
                "table_score_degradation": table_delta,
                "sensitivity_rank": None,
            }
            print(
                f"    CER={metrics['mean_cer']:.4f} "
                f"(delta={metrics['mean_cer'] - baseline['mean_cer']:+.4f})"
            )
        except Exception as exc:
            results["layer_groups"][group] = {
                "status": "error",
                "error": str(exc),
                "sensitivity_rank": None,
            }
            print(f"    ERROR: {exc}")
        finally:
            group_model = None
            group_processor = None
            gc.collect()
            mx.clear_cache()

    successful_groups = [
        (name, data)
        for name, data in results["layer_groups"].items()
        if data["status"] == "success"
    ]
    successful_groups.sort(
        key=lambda item: (
            item[1]["cer_delta"],
            item[1]["digit_cer_delta"]
            if item[1]["digit_cer_delta"] is not None
            else float("-inf"),
        ),
        reverse=True,
    )
    for rank, (_, data) in enumerate(successful_groups, start=1):
        data["sensitivity_rank"] = rank

    # Save results
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(results, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    print(f"\nResults saved to: {args.output}")
    if any(data["status"] == "error" for data in results["layer_groups"].values()):
        raise SystemExit(1)


if __name__ == "__main__":
    main()