""" Recall — shared inference wrapper. OWNER: Nikolai (Module B) Everything that touches the model goes through `chat()`. Both content_pipeline and learning_engine import this and nothing else model-related. Default is STUB mode (RECALL_STUB=1) so `python app.py` runs with no GPU and no model download. Flip RECALL_STUB=0 once the real MiniCPM call works on the Space. Model is a one-env-var config flip (NAH-9). Set RECALL_MODEL to a known alias or any full HF id; default is the 8B. If the Space is too slow / OOM, swap to a smaller model with no code change: RECALL_MODEL=1b RECALL_STUB=0 python app.py # MiniCPM5-1B (fast fallback) RECALL_MODEL=4b RECALL_STUB=0 python app.py # MiniCPM3-4B (Tiny Titan badge) Aliases resolve via MODELS below; an unknown value is treated as a literal HF id. Load dtype/device default to bf16 + device_map="auto" (correct for the Space's CUDA GPU). For a local real-model smoke test on Apple Silicon, override them — bf16 on MPS produces garbage, so use CPU/float32: RECALL_STUB=0 RECALL_MODEL=1b RECALL_DTYPE=float32 RECALL_DEVICE=cpu python app.py """ from __future__ import annotations import json import os import re STUB = os.getenv("RECALL_STUB", "1") == "1" # Known models, keyed by short alias so swapping is a single env-var flip. MODELS = { "8b": "openbmb/MiniCPM4.1-8B", # default / primary "1b": "openbmb/MiniCPM5-1B", # fast fallback if the Space is slow / OOM "4b": "openbmb/MiniCPM3-4B", # mid fallback (Tiny Titan badge) } _requested = os.getenv("RECALL_MODEL", "8b") # Accept an alias ("1b") or a full HF id ("org/model") passed through verbatim. MODEL_ID = MODELS.get(_requested, _requested) _model = None _tokenizer = None def active_model() -> str: """The HF model id currently configured ('stub' when running stubbed).""" return "stub" if STUB else MODEL_ID # Load-time dtype/device, overridable for local dev (defaults are correct for # the Space's CUDA GPU). bf16 on Apple-Silicon MPS produces garbage output, so a # Mac real-model smoke test needs RECALL_DTYPE=float32 RECALL_DEVICE=cpu; unset, # behavior is unchanged (bf16 + device_map="auto"). _DTYPE_ALIASES = { "bfloat16": "bfloat16", "bf16": "bfloat16", "float16": "float16", "fp16": "float16", "half": "float16", "float32": "float32", "fp32": "float32", "float": "float32", } def _resolve_dtype_name() -> str: """Normalized torch dtype name from RECALL_DTYPE (default 'bfloat16'). Unknown values fall back to the default rather than erroring at load.""" return _DTYPE_ALIASES.get(os.getenv("RECALL_DTYPE", "bfloat16").lower(), "bfloat16") def _resolve_device_map(): """device_map for from_pretrained. Default 'auto' (accelerate places it); RECALL_DEVICE overrides, e.g. 'cpu' for stable local CPU inference.""" return os.getenv("RECALL_DEVICE") or "auto" def _load(): """Lazy-load the model once. Only called when STUB is off.""" global _model, _tokenizer if _model is not None: return import torch from transformers import AutoModelForCausalLM, AutoTokenizer dtype = getattr(torch, _resolve_dtype_name()) device_map = _resolve_device_map() print(f"[recall] loading model: {MODEL_ID} (dtype={_resolve_dtype_name()}, " f"device_map={device_map})") _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) _model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=dtype, device_map=device_map, trust_remote_code=True, ) def _render_prompt(messages: list[dict]) -> str: """Build the prompt string. MiniCPM4.1/MiniCPM5 are hybrid reasoning models; we pass enable_thinking=False so they answer directly instead of spending the (deliberately tight) token budget on a preamble that would push the JSON answer past max_tokens — and slow the demo. Non-reasoning models (e.g. MiniCPM3-4B) ignore the unused template variable; templates that actively reject it fall back to a plain render. extract_json() still strips any that leaks through, so this is an optimization, not a correctness dependency.""" try: return _tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) except Exception: # noqa: BLE001 — template can't take the flag; render plain return _tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) def chat(messages: list[dict], max_tokens: int = 512) -> str: """ messages: [{"role": "system"|"user"|"assistant", "content": str}, ...] Returns the assistant's text. On a HF ZeroGPU Space, wrap the *callers'* entrypoints (or this function) with @spaces.GPU. Keep max_tokens tight — latency is the demo killer. """ if STUB: return _stub_reply(messages) _load() text = _render_prompt(messages) inputs = _tokenizer(text, return_tensors="pt").to(_model.device) out = _model.generate( **inputs, max_new_tokens=max_tokens, do_sample=True, temperature=0.7, top_p=0.9, ) gen = out[0][inputs["input_ids"].shape[1]:] return _tokenizer.decode(gen, skip_special_tokens=True).strip() # ---- JSON helper: model output is never trusted ---------------------------- _THINK_CLOSE = re.compile(r"", re.IGNORECASE) def _strip_think(text: str) -> str: """Drop a reasoning-model preamble. MiniCPM4.1/MiniCPM5 are hybrid reasoning models that emit … before the actual answer; when the chat template pre-fills the opening tag only the closing shows up in the reply. Either way the answer (the JSON we want) is whatever follows the LAST , so anchoring there also defuses stray braces inside the reasoning that would otherwise mislead the JSON search below. A truncated, never-closed leaves the text untouched -> extract_json returns None -> the caller's repair retry / safe default handles it.""" last = None for last in _THINK_CLOSE.finditer(text): pass return text[last.end():].strip() if last else text def _loads(s: str): """json.loads, but tolerant of models that over-escape their output. Seen with MiniCPM4.1-8B, which sometimes escapes JSON as if it were a string literal — quotes as `\\"` and newlines as `\\n` — e.g. `[\\n {\\"k\\": \\"v\\"}\\n]` instead of real JSON. If the straight parse fails and the text carries `\\"`, retry by (a) decoding it as a JSON string body, which undoes \\", \\n, \\t and unicode escapes in one shot, then parsing the result; and (b) a simpler quote-only un-escape as a backstop. Strictly additive: valid JSON parses on the first try and never reaches the fallbacks, so legitimately escaped quotes inside a string are untouched. Returns the parsed value or None.""" try: return json.loads(s) except Exception: pass if '\\"' in s: # (a) Treat the whole reply as an escaped string and decode it once. try: return json.loads(json.loads('"' + s + '"')) except Exception: pass # (b) Backstop: just collapse the escaped quotes. try: return json.loads(s.replace('\\"', '"')) except Exception: pass return None def extract_json(text: str): """ Pull the first JSON object/array out of a model reply. Returns the parsed object or None. Callers should handle None (skip card / use fallback grade). """ text = _strip_think(text.strip()) # strip ```json fences if present text = re.sub(r"^```(?:json)?|```$", "", text, flags=re.MULTILINE).strip() data = _loads(text) if data is not None: return data match = re.search(r"(\[.*\]|\{.*\})", text, re.DOTALL) if match: return _loads(match.group(1)) return None def chat_json(messages: list[dict], max_tokens: int = 256, retries: int = 1): """ Call the model and parse its reply as JSON, with up to `retries` repair passes. Model output is never trusted: if the first reply isn't valid JSON we feed it back with a terse "return ONLY valid JSON" nudge and try again. Returns the parsed object/array, or None if every attempt fails (callers must handle None with a safe default — never crash the study loop). """ convo = list(messages) for attempt in range(retries + 1): reply = chat(convo, max_tokens=max_tokens) data = extract_json(reply) if data is not None: return data if attempt < retries: # Repair pass: show the model its bad reply and demand clean JSON. convo = messages + [ {"role": "assistant", "content": reply}, {"role": "user", "content": "That was not valid JSON. Reply again with ONLY the JSON " "value — no prose, no markdown fences, no trailing text."}, ] return None # ---- Stub replies so the app runs with no model ---------------------------- def _stub_reply(messages: list[dict]) -> str: """Cheap deterministic-ish replies keyed off the caller's intent tag.""" content = " ".join(m.get("content", "") for m in messages).lower() if "generate" in content and "question" in content: return json.dumps([ {"question": "[stub] What is the main idea of the source text?", "answer": "The main concept described in the passage.", "topic": "Stub Topic", "difficulty": 1}, {"question": "[stub] How does the key concept apply in this context?", "answer": "It applies by connecting the described mechanism to the outcome.", "topic": "Stub Topic", "difficulty": 2}, {"question": "[stub] Compare and contrast the two ideas presented.", "answer": "They differ in scope but share the same underlying principle.", "topic": "Stub Topic", "difficulty": 3}, ]) if "grade" in content or "score" in content: return json.dumps({ "score": 4, "explanation": "[stub] Close — you captured the main idea but missed a detail.", "missed_concept": "the specific detail", }) if "follow" in content: return json.dumps([ {"question": "[stub follow-up] Can you restate the missed detail?", "answer": "The specific detail from the passage.", "topic": "Stub Topic"}, ]) return "[stub] model reply"