llm.py

"""BioMistral-7B loader + ZeroGPU-decorated inference for the TuringDNA assistant.

Architecture decisions worth pinning here so the next person reading this
doesn't waste a day re-deriving them:

1.  bf16 not bnb-4bit.
    ZeroGPU shares an NVIDIA A100 40 GB. A 7-B model in bf16 is ~14 GB —
    fits with 26 GB to spare. The original spec called for bitsandbytes
    4-bit, but that adds dequant overhead on every forward pass for no
    memory benefit when we already have 26 GB free. 4-bit is the right
    call only when memory is the constraint.

2.  Model loaded at import, moved to GPU inside @spaces.GPU.
    ZeroGPU's pattern: load on CPU once (slow first cold-start, then
    cached on the host), then move to GPU only during a decorated
    function call. The function "borrows" the GPU for up to its declared
    duration, then ZeroGPU reclaims it for the next caller. This is how
    one A100 serves dozens of Spaces.

3.  No streaming over the gradio_client boundary (yet).
    Gradio supports streaming for its own UI via TextIteratorStreamer
    and yielding from the chat_fn. But gradio_client.predict() returns
    the full string at the end of generation; intermediate chunks aren't
    surfaced. Adding streaming requires either WebSocket plumbing or
    Server-Sent Events on a custom FastAPI route. MVP returns the full
    response in one shot.

4.  Fallback model.
    BioMistral-7B is gated/can be removed. If load fails, we fall back to
    mistralai/Mistral-7B-Instruct-v0.2 (Apache-2.0, always available)
    rather than crashing the Space.

5.  Mistral instruct template, not chat-template.
    BioMistral inherits Mistral's [INST]…[/INST] template. Transformers'
    tokenizer.apply_chat_template() does the right thing if the model
    config carries a chat_template — for BioMistral-7B it doesn't, so we
    build the template manually below.
"""

from __future__ import annotations

import logging
from typing import Iterable, List, Optional

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

logger = logging.getLogger("turingdna.assistant")
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")

# ---------------------------------------------------------------------- spaces
# `spaces` is the HF SDK that provides @spaces.GPU. On a non-ZeroGPU
# host (local dev, plain CPU Space) the import will succeed but the
# decorator becomes a no-op. Wrap the import so local pip-only
# environments don't crash either.
try:
    import spaces  # type: ignore[import-not-found]
    HAS_SPACES = True
except ImportError:  # pragma: no cover — local dev path
    HAS_SPACES = False

    class _SpacesShim:
        @staticmethod
        def GPU(*decorator_args, **decorator_kwargs):  # noqa: N802 — match upstream
            """No-op decorator stand-in for environments without spaces SDK."""
            # spaces.GPU is used either as @spaces.GPU or @spaces.GPU(duration=N).
            # Detect which calling convention was used and handle both.
            if decorator_args and callable(decorator_args[0]) and not decorator_kwargs:
                return decorator_args[0]
            def _wrap(fn):
                return fn
            return _wrap

    spaces = _SpacesShim()  # type: ignore[assignment]


# ─────────────────────────────────────────────────────────── model loading

PRIMARY_MODEL = "BioMistral/BioMistral-7B"
FALLBACK_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"

GENERATION_DEFAULTS = {
    "max_new_tokens": 512,
    "temperature": 0.3,
    "top_p": 0.92,
    "repetition_penalty": 1.05,
    "do_sample": True,
}

SYSTEM_PROMPT = """You are a protein biology expert helping a researcher use TuringDNA, a directed-evolution engine.

ANSWER DIRECTLY. If the question has a one-word answer (e.g. "what's the start codon in E. coli?" → "ATG (DNA) / AUG (mRNA)."), give that answer with one short follow-up sentence of reasoning. Don't ask for clarification on questions that aren't ambiguous.

NEVER introduce yourself. Don't say "Hey there! I'm BioMistral…" Don't say "I'm here to help you with…" Don't recite your identity facts. Just answer the question.

Greetings: if the user says "hi", "hello", "hey", "yo", or any plain greeting with no question — respond with a single short line like "Hi — what are you working on?" or "Hello. Ask away." DO NOT introduce yourself.

Short follow-ups: if the user message is short (under 10 words) and is clearly a follow-up to the previous turn (e.g. "yes", "no", "cDNA then", "what about yeast", "and why"), interpret it as continuing the previous topic. The prior exchange is usually included as context above the current message — use it.

Voice: warm, peer-to-peer, first person. Short paragraphs. Bullets only when the answer is genuinely enumerable. No "Let me know if you have more questions!" tail.

Domain: enzyme mechanisms, ESM-2 ΔLL scoring (POSITIVE ΔLL = more likely than WT under the model = predicted to be TOLERATED; NEGATIVE = predicted DISRUPTIVE), directed evolution strategy, codon optimization (E. coli / yeast / human), cloning (Golden Gate / Gibson / restriction-ligation, common vectors).

Honesty: don't invent active-site residues, domain boundaries, or PDB/UniProt IDs you're not sure about — say "I'd check UniProt for that." ΔLL is a likelihood under a language model, not a measurement of function — don't promise activity, just predicted tolerance.

If — and ONLY if — directly asked "who are you" or "who created you" or "what model is this": say "BioMistral, an open-source biomedical Mistral-7B fine-tune. TuringDNA deploys me here." One sentence. Don't volunteer this on any other question."""


def _load(model_id: str) -> tuple[AutoModelForCausalLM, AutoTokenizer]:
    """Tokenizer + model in bf16 on CPU. Move-to-cuda happens later inside @spaces.GPU."""
    logger.info("Loading tokenizer for %s …", model_id)
    tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    if tok.pad_token_id is None:
        tok.pad_token_id = tok.eos_token_id

    logger.info("Loading model weights for %s (bf16) …", model_id)
    mdl = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        # Do NOT set device_map="auto" — ZeroGPU manages placement
        # explicitly inside the decorated function. device_map="auto"
        # tries to pin to "cuda" at load time, which on ZeroGPU's host
        # CPU has no CUDA visible and crashes.
        low_cpu_mem_usage=True,
    )
    mdl.eval()
    logger.info("Loaded %s — %.2f GB on CPU.", model_id,
                sum(p.numel() * p.element_size() for p in mdl.parameters()) / 1e9)
    return mdl, tok


def _load_with_fallback() -> tuple[AutoModelForCausalLM, AutoTokenizer, str]:
    """Try BioMistral first; fall back to vanilla Mistral if it can't be fetched.
    Returns (model, tokenizer, model_id_actually_used)."""
    for model_id in (PRIMARY_MODEL, FALLBACK_MODEL):
        try:
            mdl, tok = _load(model_id)
            return mdl, tok, model_id
        except Exception as exc:  # noqa: BLE001
            logger.warning("Could not load %s: %s", model_id, exc)
            if model_id == FALLBACK_MODEL:
                raise
    raise RuntimeError("Unreachable: both primary and fallback failed.")


# Loaded once per Space process lifetime. Cold-start hits this; subsequent
# wakeups skip the download (HF caches in /data on the Space) but still
# pay the ~10-15 s deserialize cost.
MODEL, TOKENIZER, MODEL_ID_IN_USE = _load_with_fallback()


# ────────────────────────────────────────────────────── prompt formatting

def format_mistral_prompt(
    user_message: str,
    history: Optional[List[dict]] = None,
    system_prompt: str = SYSTEM_PROMPT,
    max_history_turns: int = 8,
) -> str:
    """Build a Mistral instruct-template prompt from a system message,
    an optional history of {role, content} dicts (OpenAI-style), and a
    new user message.

    Format reference (Mistral-7B-Instruct-v0.1+):
        <s>[INST] {sys}\\n\\n{u1} [/INST]{a1}</s>[INST] {u2} [/INST]{a2}</s>[INST] {u3} [/INST]

    System prompt is prepended INSIDE the first [INST] block — Mistral
    has no dedicated system role token. Trailing assistant span is left
    open for generation (no </s>).

    History is truncated to the last ``max_history_turns`` user/assistant
    pairs so the prompt doesn't grow unbounded. Older context is dropped
    silently; if you need long-term memory, summarize before passing.
    """
    history = list(history or [])

    # Coerce to a list of (user, assistant) pairs. The shape of `history`
    # depends on how the ChatInterface is called:
    #   - Via the Gradio UI with type="messages":
    #       [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ...]
    #   - Via the Gradio UI with type="tuples" (legacy):
    #       [["user msg", "assistant msg"], ...]
    #   - Via gradio_client API: history may be empty OR may arrive in
    #     either shape depending on the session state — Gradio 4.44 is
    #     inconsistent here. Earlier this code crashed with
    #         AttributeError: 'list' object has no attribute 'get'
    #     because msg arrived as a [user, assistant] tuple, not a dict.
    # Handle both shapes defensively.
    pairs: list[tuple[str, str]] = []
    pending_user: Optional[str] = None
    for msg in history:
        # Dict (modern messages format)
        if isinstance(msg, dict):
            role = (msg.get("role") or "").lower()
            content = (msg.get("content") or "").strip()
            if not content:
                continue
            if role == "user":
                pending_user = content
            elif role == "assistant" and pending_user is not None:
                pairs.append((pending_user, content))
                pending_user = None
        # List/tuple (legacy tuples format: [user_msg, assistant_msg])
        elif isinstance(msg, (list, tuple)) and len(msg) >= 2:
            u = str(msg[0] or "").strip()
            a = str(msg[1] or "").strip()
            if u and a:
                pairs.append((u, a))
        # Anything else — skip, don't crash
        else:
            continue
    # Drop any dangling user message without a paired assistant —
    # the new user_message at the end already covers "what's the user
    # asking now"; dangling halves would just duplicate it.

    # Truncate.
    if max_history_turns > 0 and len(pairs) > max_history_turns:
        pairs = pairs[-max_history_turns:]

    # Build the prompt. Mistral expects <s> at the very start (the
    # tokenizer adds it via add_special_tokens=True, but we include it
    # explicitly so the format is self-evident in logs).
    parts: list[str] = ["<s>"]
    # First [INST] carries the system prompt prepended.
    if pairs:
        first_u, first_a = pairs[0]
        parts.append(f"[INST] {system_prompt}\n\n{first_u} [/INST]{first_a}</s>")
        for u, a in pairs[1:]:
            parts.append(f"[INST] {u} [/INST]{a}</s>")
        parts.append(f"[INST] {user_message} [/INST]")
    else:
        # No history — system prompt + new message in the first INST.
        parts.append(f"[INST] {system_prompt}\n\n{user_message} [/INST]")

    return "".join(parts)


# ──────────────────────────────────────────────────────────── inference

@spaces.GPU(duration=30)
def generate(prompt: str, **gen_kwargs) -> str:
    """Run one generation pass and return only the newly-generated text.

    Decorated with @spaces.GPU(duration=30) so ZeroGPU acquires the A100
    for up to 30 seconds. ZeroGPU bills the FULL declared duration
    against the caller's daily quota regardless of how long the actual
    call took — so a 60s budget for a 5s response wastes 55s of quota.
    Most chat answers complete in 5-15s; 30s is comfortably above that
    while halving the quota cost vs. the old 60s. If responses ever
    start truncating because they hit the timeout, bump back up.

    On a local CPU host the decorator is a no-op (see the spaces shim
    at the top of this file) and the call runs on whatever device the
    model currently sits on.
    """
    cfg = {**GENERATION_DEFAULTS, **gen_kwargs}

    # Move to CUDA if available. On ZeroGPU this is the line that
    # acquires the actual GPU; on a local CPU host it's a no-op since
    # torch.cuda.is_available() returns False.
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if next(MODEL.parameters()).device.type != device:
        MODEL.to(device)

    inputs = TOKENIZER(prompt, return_tensors="pt", truncation=True, max_length=4096).to(device)
    input_token_count = inputs.input_ids.shape[1]

    with torch.no_grad():
        outputs = MODEL.generate(
            **inputs,
            pad_token_id=TOKENIZER.pad_token_id,
            eos_token_id=TOKENIZER.eos_token_id,
            **cfg,
        )

    # Skip the prompt tokens — only return the assistant's new tokens.
    new_tokens = outputs[0][input_token_count:]
    text = TOKENIZER.decode(new_tokens, skip_special_tokens=True).strip()
    return text


def chat(user_message: str, history: Optional[List[dict]] = None, **gen_kwargs) -> str:
    """Top-level convenience: takes a chat message + history, returns the
    assistant's reply. This is what Gradio's ChatInterface calls."""
    prompt = format_mistral_prompt(user_message, history)
    return generate(prompt, **gen_kwargs)