Spaces:
Running
Running
| """Real checkpoint inference, with bounded logit and attention storage.""" | |
| from __future__ import annotations | |
| import os | |
| import threading | |
| import time | |
| import uuid | |
| from collections import OrderedDict | |
| import numpy as np | |
| import torch | |
| from jinja2 import TemplateError | |
| from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer | |
| MODEL_ID = "Qwen/Qwen3-0.6B" | |
| REVISION = "c1899de289a04d12100db370d81485cdf75e47ca" | |
| GEMMA_ID = "ThomasTheMaker/gm3-270m-TinyGSM-reasoning" | |
| MODELS = { | |
| MODEL_ID: dict(name="Qwen3-0.6B", revision=REVISION, native_thinking=True), | |
| GEMMA_ID: dict(name="Gemma 3 270M · TinyGSM reasoning", | |
| revision="aa034f2832467400d4e0a09df9fae0c517d41579", native_thinking=False), | |
| } | |
| MAX_TOKENS = int(os.environ.get("QWEN_MAX_TOKENS", "512")) | |
| class Runtime: | |
| def __init__(self, model_id=MODEL_ID): | |
| self.model_id = model_id | |
| self.spec = MODELS[model_id] | |
| self.config = None | |
| self.phase = "loading" | |
| self.message = f"Loading {self.spec['name']}…" | |
| self.tokenizer = None | |
| self.model = None | |
| self.device = "cpu" | |
| self.lock = threading.Lock() | |
| self.tokenizer_lock = threading.Lock() | |
| self.sessions = OrderedDict() | |
| self.active = None | |
| def load(self): | |
| try: | |
| torch.set_num_threads(min(6, os.cpu_count() or 4)) | |
| self.config = AutoConfig.from_pretrained(self.model_id, revision=self.spec["revision"]) | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, revision=self.spec["revision"]) | |
| self.message = f"Loading the {self.spec['name']} checkpoint…" | |
| requested = os.environ.get("QWEN_DEVICE", "auto") | |
| self.device = ("mps" if torch.backends.mps.is_available() else "cpu") if requested == "auto" else requested | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| self.model_id, revision=self.spec["revision"], config=self.config, dtype=torch.float32, | |
| attn_implementation="eager", | |
| ).to(self.device).eval() | |
| self.phase = "ready" | |
| self.message = f"{self.spec['name']} · {self.device.upper()} · float32 · eager attention" | |
| except Exception as exc: | |
| self.phase = "error" | |
| self.message = f"Could not load the model: {exc}" | |
| import traceback | |
| traceback.print_exc() | |
| def metadata(self): | |
| config = self.config | |
| return dict(model=self.model_id, name=self.spec["name"], revision=self.spec["revision"], | |
| native_thinking=self.spec["native_thinking"], | |
| layers=config.num_hidden_layers if config else None, | |
| heads=config.num_attention_heads if config else None, | |
| kv_heads=config.num_key_value_heads if config else None, | |
| vocab_size=config.vocab_size if config else None, | |
| attention_scale=getattr(config, "query_pre_attn_scalar", config.head_dim) if config else None, | |
| layer_types=getattr(config, "layer_types", None) if config else None, | |
| sliding_window=getattr(config, "sliding_window", None) if config else None) | |
| def status(self): | |
| result = dict(phase=self.phase, message=self.message, tokenizer_ready=self.tokenizer is not None, | |
| device=self.device, max_tokens=MAX_TOKENS, active=self.active, **self.metadata()) | |
| if self.tokenizer: | |
| result["special_tokens"] = [dict(id=i, text=t.content, special=t.special) | |
| for i, t in self.tokenizer.added_tokens_decoder.items() | |
| if i < self.config.vocab_size and not t.content.startswith("<unused") | |
| and (t.special or t.content.startswith(("<", "[")))] | |
| return result | |
| def prepare(self, payload): | |
| if not self.tokenizer: | |
| raise ValueError("The tokenizer is still loading. Try again in a few seconds.") | |
| if payload.mode == "chat": | |
| if payload.thinking and not self.spec["native_thinking"]: | |
| raise ValueError("This Gemma checkpoint has no native reasoning switch. It generates learned Reasoning / Solution text.") | |
| if not payload.messages or not any(m.role == "user" for m in payload.messages): | |
| raise ValueError("Add at least one user message.") | |
| if payload.messages[-1].role != "user": | |
| raise ValueError("The last message must have the user role before a new response. Use Raw for an arbitrary prefix.") | |
| kwargs = {"enable_thinking": payload.thinking} if self.spec["native_thinking"] else {} | |
| try: | |
| text = self.tokenizer.apply_chat_template( | |
| [m.model_dump() for m in payload.messages], tokenize=False, | |
| add_generation_prompt=True, **kwargs, | |
| ) | |
| except TemplateError as exc: | |
| raise ValueError(f"Invalid conversation for this model's chat template: {exc}") from exc | |
| else: | |
| text = payload.text | |
| with self.tokenizer_lock: | |
| ids = self.tokenizer.encode(text, add_special_tokens=False) | |
| if not ids: | |
| raise ValueError("The context is empty. Enter text or choose a template.") | |
| if any(i >= self.config.vocab_size for i in ids): | |
| raise ValueError("The input contains a tokenizer marker outside this text checkpoint's vocabulary. Remove unsupported image or other modality markers.") | |
| return text, ids | |
| def token(self, token_id): | |
| tid = int(token_id) | |
| return dict(id=tid, text=self.tokenizer.decode([tid], clean_up_tokenization_spaces=False), | |
| piece=self.tokenizer.convert_ids_to_tokens(tid), special=tid in self.tokenizer.all_special_ids) | |
| def append_ids(self, sid, text=None): | |
| """Append one real token; the existing context IDs must never change.""" | |
| previous = self.session(sid)["response"] | |
| if text is None: | |
| return [previous["predictions"][-1]["ids"][-1]] | |
| with self.tokenizer_lock: | |
| ids = self.tokenizer.encode(text, add_special_tokens=False)[:1] | |
| if not ids: | |
| raise ValueError("Enter a token or text to append. Spaces are preserved.") | |
| if any(tid >= self.config.vocab_size for tid in ids): | |
| raise ValueError("The text contains a tokenizer marker outside this text checkpoint's vocabulary.") | |
| return ids | |
| def analyze(self, payload, prefix=None, *, input_count=None, token_sources=None): | |
| if self.phase != "ready": | |
| raise ValueError(self.message) | |
| if not self.lock.acquire(blocking=False): | |
| raise BlockingIOError("The model is already computing. Wait for the current run to finish.") | |
| started = time.perf_counter() | |
| try: | |
| text, ids = self.prepare(payload) if prefix is None else prefix | |
| prefix_count = len(ids) | |
| if input_count is None: | |
| input_count = prefix_count | |
| sources = list(token_sources) if token_sources is not None else ["input"] * prefix_count | |
| if prefix_count + payload.max_new_tokens > MAX_TOKENS: | |
| raise ValueError(f"Received {prefix_count} context tokens + {payload.max_new_tokens} continuation tokens. The inspector limit is {MAX_TOKENS}; input has not been truncated. Shorten the context or continuation.") | |
| forward_message = f"Forward pass: {self.config.num_hidden_layers} layers × {self.config.num_attention_heads} heads…" | |
| self.active = "Generating the continuation…" if payload.max_new_tokens else forward_message | |
| input_ids = torch.tensor([ids], device=self.device) | |
| if payload.max_new_tokens: | |
| generated = self.model.generate( | |
| input_ids, attention_mask=torch.ones_like(input_ids), | |
| max_new_tokens=payload.max_new_tokens, do_sample=False, | |
| pad_token_id=self.tokenizer.eos_token_id, use_cache=True, | |
| ) | |
| ids = generated[0].tolist() | |
| input_ids = generated | |
| sources.extend(["predicted"] * (len(ids) - prefix_count)) | |
| self.active = forward_message | |
| output = self.model.model(input_ids, output_hidden_states=True, output_attentions=True, use_cache=False) | |
| if output.attentions is None or any(a is None for a in output.attentions): | |
| raise RuntimeError("The model did not return attention. Eager attention is required.") | |
| # Float16 is only used to retain attention maps. Computation and lens are float32. | |
| # Transfer float32 first, then cast in NumPy into independently owned memory. | |
| # Fusing MPS→CPU and float16 conversion can corrupt strided attention tensors. | |
| attention = [a[0].cpu().contiguous().numpy().astype(np.float16, copy=True) | |
| for a in output.attentions] | |
| for matrix in attention: | |
| sums = matrix.astype(np.float32).sum(axis=-1) | |
| if (not np.isfinite(matrix).all() or not np.allclose(sums, 1, atol=1e-3) | |
| or np.count_nonzero(np.triu(matrix, 1))): | |
| raise RuntimeError("An attention map failed the normalization or causal mask check.") | |
| top_ids, top_probs = [], [] | |
| for layer in range(self.model.config.num_hidden_layers): | |
| self.active = f"Logit lens: layer {layer + 1} / {self.config.num_hidden_layers}…" | |
| # HF returns the last hidden state after final RMSNorm; do not normalize twice. | |
| hidden = output.hidden_states[layer + 1] | |
| if layer != self.model.config.num_hidden_layers - 1: | |
| hidden = self.model.model.norm(hidden) | |
| layer_ids, layer_probs = [], [] | |
| for chunk in hidden[0].split(32): | |
| logits = self.model.lm_head(chunk).float() | |
| softcap = getattr(self.config, "final_logit_softcapping", None) | |
| if softcap is not None: | |
| logits = torch.tanh(logits / softcap) * softcap | |
| values, indices = logits.topk(5, dim=-1) | |
| # Center first: early Gemma states can produce large logits, where | |
| # subtracting an uncentered logsumexp loses float32 precision. | |
| maximum = values[:, :1] | |
| probabilities = (values - maximum - (logits - maximum).logsumexp(dim=-1, keepdim=True)).exp() | |
| layer_ids.extend(indices.cpu().tolist()) | |
| layer_probs.extend(probabilities.cpu().tolist()) | |
| top_ids.append(layer_ids) | |
| top_probs.append(layer_probs) | |
| del output | |
| sid = uuid.uuid4().hex | |
| vocab_ids = {tid for layer in top_ids for row in layer for tid in row} | |
| token_dict = {str(tid): self.token(tid) for tid in vocab_ids} | |
| response = dict( | |
| id=sid, context=text, tokens=[self.token(tid) for tid in ids], input_count=input_count, | |
| **self.metadata(), device=self.device, max_tokens=MAX_TOKENS, | |
| predictions=[dict(ids=[r[0] for r in layer], probabilities=[r[0] for r in probs]) | |
| for layer, probs in zip(top_ids, top_probs)], | |
| vocabulary=token_dict, elapsed=round(time.perf_counter() - started, 2), | |
| token_sources=sources, added_count=len(ids) - input_count, | |
| continuation_text=self.tokenizer.decode(ids[input_count:], clean_up_tokenization_spaces=False), | |
| generated_text=self.tokenizer.decode([tid for tid, source in zip(ids, sources) if source == "predicted"], clean_up_tokenization_spaces=False), | |
| generated_count=sources.count("predicted"), manual_count=sources.count("manual"), | |
| stopped_on_eos=bool(len(ids) > input_count and ids[-1] == self.tokenizer.eos_token_id), | |
| stop_token=self.token(ids[-1])["text"] if len(ids) > input_count else None, | |
| thinking=payload.thinking if payload.mode == "chat" and self.spec["native_thinking"] else None, | |
| ) | |
| self.sessions[sid] = dict(attention=attention, top_ids=top_ids, top_probs=top_probs, | |
| response=response) | |
| # Bound memory; even at 512 tokens, at most ~470 MB of attention is retained. | |
| while len(self.sessions) > 2: | |
| self.sessions.popitem(last=False) | |
| return response | |
| finally: | |
| self.active = None | |
| self.lock.release() | |
| def session(self, sid): | |
| session = self.sessions.get(sid) | |
| if session is None: | |
| raise KeyError("This result has been evicted from memory. Run the analysis again.") | |
| return session | |
| def inspect(self, sid, layer, query): | |
| s = self.session(sid) | |
| weights = s["attention"][layer][:, query, :].astype(np.float32) | |
| return dict(top=[dict(**self.token(tid), probability=p) for tid, p in | |
| zip(s["top_ids"][layer][query], s["top_probs"][layer][query])], | |
| attention=weights.tolist()) | |
| def attention_layers(self, sid, query, head="mean"): | |
| """Read one query across layers, without another model forward pass.""" | |
| session = self.session(sid) | |
| weights = [] | |
| for matrix in session["attention"]: | |
| if head == "mean": | |
| row = matrix[:, query, :].astype(np.float32).mean(axis=0) | |
| else: | |
| row = matrix[head, query, :].astype(np.float32) | |
| weights.append(row.tolist()) | |
| return dict(query=query, head=head, weights=weights) | |