llm.py

"""
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 <think> 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
    <think> 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"</think\s*>", re.IGNORECASE)


def _strip_think(text: str) -> str:
    """Drop a reasoning-model <think> preamble. MiniCPM4.1/MiniCPM5 are hybrid
    reasoning models that emit <think>…</think> before the actual answer; when
    the chat template pre-fills the opening tag only the closing </think> shows
    up in the reply. Either way the answer (the JSON we want) is whatever follows
    the LAST </think>, so anchoring there also defuses stray braces inside the
    reasoning that would otherwise mislead the JSON search below. A truncated,
    never-closed <think> 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"