Blame for runtime.py

"""Live inference using the supplied LLMON parser and attention extractor."""

import json
import logging
import os
import re
from pathlib import Path

from pydantic import BaseModel, ConfigDict, Field, model_validator

from vendor import span_parser

ROOT = Path(__file__).resolve().parent
DEFAULT_MODEL_REPO = "ibm-granite/granite-4.1-3b-llmon-preview"
LEGACY_MODEL_REPO = "ibm-research/granite-4.1-3b-llmon-preview"
MODEL_ID = os.getenv("LLMON_MODEL_ID", "g413b-lmon-exp-207-x")
MODEL_LABEL = os.getenv("LLMON_MODEL_LABEL", "Granite LLMON")
logger = logging.getLogger("uvicorn.error")
EXAMPLES = {
    row["id"]: row
    for line in (ROOT / "new_dlb_examples.jsonl").read_text().splitlines()
    if line.strip() and (row := json.loads(line))
}


def configured_model_source() -> tuple[str, str]:
    """Select the owner's new repository, including Spaces with old variables."""
    repo = os.getenv("LLMON_MODEL_REPO", "").strip() or DEFAULT_MODEL_REPO
    revision = os.getenv("LLMON_MODEL_REVISION", "").strip() or "main"
    if repo == LEGACY_MODEL_REPO:
        # A commit in the old repository is not a verified commit in the new one.
        logger.warning(
            "Migrating LLMON_MODEL_REPO from %s to %s at main; the old revision is not reused. "
            "Set the new repository and pin its own revision in Space settings for reproducibility.",
            LEGACY_MODEL_REPO,
            DEFAULT_MODEL_REPO,
        )
        return DEFAULT_MODEL_REPO, "main"
    return repo, revision


class AttentionRequest(BaseModel):
    """Bound a live request to the configured model and supported prompt format."""

    model_config = ConfigDict(extra="forbid")
    model_id: str = MODEL_ID
    example_id: int | None = Field(default=None, ge=0)
    prompt_llmon: str | None = Field(default=None, min_length=1, max_length=16384)
    dataset_tag: str | None = None
    max_tokens: int = Field(default=64, ge=1, le=64)
    layer: int | None = Field(default=None, ge=0, le=255)

    @model_validator(mode="after")
    def validate_selection(self) -> "AttentionRequest":
        """Reject ambiguous requests and unknown model/dataset identifiers."""
        if self.model_id != MODEL_ID:
            raise ValueError("Select the configured LLMON model.")
        if self.dataset_tag not in (None, "default"):
            raise ValueError("Uploaded datasets are sent as individual prompts.")
        if (self.example_id is None) == (self.prompt_llmon is None):
            raise ValueError("Provide either an example ID or one LLMON prompt.")
        if self.example_id is not None and self.example_id not in EXAMPLES:
            raise ValueError("Unknown example ID.")
        return self

    def prompt(self) -> str:
        """Resolve a prompt without reading any user-selected filesystem path."""
        return self.prompt_llmon if self.prompt_llmon is not None else EXAMPLES[self.example_id]["llmon_prompt_dirty"]


class InferenceEngine:
    """Keep one configured checkpoint in memory, with no fallback recordings."""

    def __init__(self) -> None:
        """Initialize without downloading models or requiring CUDA."""
        self.model = None
        self.tokenizer = None
        self.model_info = {}
        self.status = "not_configured"
        self.message = "The live model is not configured yet."
        self.max_input_tokens = int(os.getenv("LLMON_MAX_INPUT_TOKENS", "1024"))

    def prepare(self) -> None:
        """Load weights at startup, after spaces has enabled CUDA emulation."""
        model_path = os.getenv("LLMON_MODEL_PATH")
        import torch
        from huggingface_hub import snapshot_download

        from vendor.attention_extractor import load_model

        device = os.getenv("LLMON_DEVICE", "cuda")
        if device not in {"cuda", "cpu", "mps"}:
            raise ValueError("LLMON_DEVICE must be cuda, cpu, or mps.")
        if os.getenv("SPACE_ID") and device != "cuda":
            raise ValueError("The live Space requires CUDA; CPU fallback is disabled.")
        if device == "cuda" and not torch.cuda.is_available():
            self.status = "gpu_unavailable"
            self.message = "A GPU has not been configured for this live demo."
            return
        self.status = "loading"
        repo = None
        requested_revision = None
        resolved_revision = None
        if not model_path:
            repo, requested_revision = configured_model_source()
            model_path = snapshot_download(
                repo_id=repo,
                revision=requested_revision,
                allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt", "*.tiktoken"],
            )
            # snapshot_download returns .../snapshots/<commit>; expose that SHA,
            # not the possibly moving branch name, and never a local model path.
            snapshot_name = Path(model_path).name
            if re.fullmatch(r"[0-9a-f]{40}", snapshot_name):
                resolved_revision = snapshot_name
        self.model, self.tokenizer, self.model_info = load_model(model_path)
        self.model.to(device)
        self.model_info.update(
            model_id=MODEL_ID,
            device=device,
            model_repo=repo,
            model_revision=resolved_revision,
            model_requested_revision=requested_revision,
        )
        self.status = "ready"
        self.message = "Live model ready."
        if repo:
            logger.info("Loaded LLMON repository=%s revision=%s", repo, resolved_revision or requested_revision)

    def ensure_ready(self) -> None:
        """Stop requests when the real checkpoint is unavailable."""
        if self.status != "ready":
            raise RuntimeError(self.message)

    def parse(self, request: AttentionRequest) -> span_parser.ParseResult:
        """Check the token budget and the global-close grammar before inference."""
        self.ensure_ready()
        prompt = request.prompt()
        encoded = self.tokenizer(prompt, add_special_tokens=False)
        if len(encoded["input_ids"]) > self.max_input_tokens:
            raise ValueError(f"Prompt exceeds the {self.max_input_tokens}-token input limit.")
        parsed = span_parser.parse_llmon_prompt(prompt, self.tokenizer)
        target_instructions = [s for s in parsed.spans if s.type == "instr" and s.is_target]
        if len(target_instructions) != 1:
            raise ValueError("Use global-close LLMON syntax with one selected instruction.")
        return parsed

    def infer(self, operation: str, payload: dict) -> dict:
        """Compute actual attention with the original float32 eager implementation."""
        from vendor.attention_extractor import (
            extract_generation_attention,
            extract_prefill_attention,
        )

        request = AttentionRequest.model_validate(payload)
        parsed = self.parse(request)
        spans = [s.to_dict() for s in parsed.spans]
        base = {
            "model": self.model_info,
            "example_id": request.example_id if request.example_id is not None else -1,
            "spans": spans,
        }
        args = (self.model, self.tokenizer, request.prompt(), spans, parsed.span_labels, self.model_info)
        if operation == "prefill":
            return {
                **base,
                "tokens": parsed.tokens,
                "span_labels": parsed.span_labels,
                "attention_by_layer": extract_prefill_attention(*args),
            }
        attention_layers = [i for i, kind in enumerate(self.model_info["layer_types"]) if kind == "attention"]
        if request.layer is not None and request.layer not in attention_layers:
            raise ValueError("Select a layer with available attention data.")
        return {**base, **extract_generation_attention(*args, max_tokens=request.max_tokens, layer=request.layer)}

    def benchmark(self) -> dict:
        """Compute the three bundled examples using the upstream focus thresholds."""
        rows = []
        for example_id in EXAMPLES:
            prefill = self.infer("prefill", {"model_id": MODEL_ID, "example_id": example_id})
            layers = []
            by_layer = prefill["attention_by_layer"]
            for key in sorted(by_layer, key=int):
                layer = by_layer[key]
                layers.append(
                    {
                        "layer_index": layer["layer_index"],
                        "layer_type": layer["layer_type"],
                        **(layer.get("metrics") or {"ifr": None, "das": None, "leak": None}),
                    }
                )
            final = [layer for layer in layers if layer["layer_type"] == "attention"][-1]
            ifr, das = final["ifr"], final["das"]
            result = "success" if ifr > 0.6 and das > 0.5 else "failure" if ifr <= 0.6 and das <= 0.5 else "partial"
            diag = "correct_attn" if result == "success" else "wrong_instr" if ifr <= 0.6 else "data_leakage"
            rows.append(
                {
                    "model_id": MODEL_ID,
                    "example_id": example_id,
                    "result": result,
                    "diag_type": diag,
                    "diag_description": "Attention focus heuristic; not answer correctness.",
                    "layers": layers,
                    "final_metrics": {k: final[k] for k in ["ifr", "das", "leak"]},
                    "kv_group_specialization": by_layer[str(final["layer_index"])].get("kv_groups", []),
                }
            )
        return {"results": rows, "progress": None, "dataset_tag": "default"}


engine = InferenceEngine()