Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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() | |