"""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/; 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()