src/demo/model_loader.py

"""Model loading for Sangue e Grafi — multi-model support.

Supports Gemma 4B and Nemotron Nano 4B with LoRA adapters.
Two-step loading: SFT adapter merged into base, then GRPO adapter on top.
Models are loaded on demand and cached; only one model lives on GPU at a time.
"""

from __future__ import annotations

import gc
import logging
import os
from threading import Lock
from typing import Callable

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Model configurations
# ---------------------------------------------------------------------------

MODEL_CONFIGS = {
    "gemma": {
        "label": "Gemma 4B",
        "base_model": "google/gemma-4-E2B-it",
        "auto_class": "AutoModelForImageTextToText",
        "sft_adapter": "cyberandy/sangue-e-grafi-gemma4-e2b-sft-adversarial-v7",
        "grpo_adapter": "cyberandy/sangue-e-grafi-gemma4-e2b-grpo-run-f-v7",
        "chat_style": "gemma",  # user/model roles
        "needs_float8_patch": True,
        "trust_remote_code": False,
    },
    "nemotron": {
        "label": "Nemotron 4B",
        "base_model": "nvidia/Nemotron-Mini-4B-Instruct",
        "auto_class": "AutoModelForCausalLM",
        "sft_adapter": "cyberandy/sangue-e-grafi-nemotron-nano-sft-v7",
        "grpo_adapter": "cyberandy/sangue-e-grafi-nemotron-nano-grpo",
        "chat_style": "chatml",  # standard system/user/assistant
        "needs_float8_patch": False,
        "trust_remote_code": True,
    },
}

# Cache: model_key -> (model, tokenizer)
_cache: dict[str, tuple] = {}
_lock = Lock()


def load_model(model_key: str = "gemma"):
    """Load base + SFT (merged) + GRPO adapter. Thread-safe, cached."""
    global _cache

    with _lock:
        if model_key in _cache:
            return _cache[model_key]

        # Evict any other model from GPU — can't hold two 4B models
        for old_key in list(_cache.keys()):
            if old_key != model_key:
                logger.info(f"Evicting {old_key} model to free GPU memory")
                old_model, _ = _cache.pop(old_key)
                del old_model
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                gc.collect()

        import torch
        from transformers import AutoTokenizer
        from peft import PeftModel

        config = MODEL_CONFIGS[model_key]

        # Monkey-patch for Gemma compatibility
        if config["needs_float8_patch"] and not hasattr(torch, 'float8_e8m0fnu'):
            torch.float8_e8m0fnu = torch.float8_e4m3fn

        hf_token = os.environ.get("HF_TOKEN", "") or None
        base_id = config["base_model"]
        trust = config["trust_remote_code"]

        logger.info(f"Loading {config['label']} base: {base_id}")

        tokenizer = AutoTokenizer.from_pretrained(
            base_id, token=hf_token, trust_remote_code=trust
        )

        # Import the right auto class
        if config["auto_class"] == "AutoModelForImageTextToText":
            from transformers import AutoModelForImageTextToText as ModelClass
        else:
            from transformers import AutoModelForCausalLM as ModelClass

        model = ModelClass.from_pretrained(
            base_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            token=hf_token,
            trust_remote_code=trust,
        )
        logger.info("Base model loaded")

        # Step 1: Load and merge SFT adapter into base
        sft_id = config["sft_adapter"]
        if sft_id:
            logger.info(f"Loading SFT adapter: {sft_id}")
            try:
                model = PeftModel.from_pretrained(model, sft_id, token=hf_token)
                model = model.merge_and_unload()
                logger.info("SFT adapter merged into base")
            except Exception as e:
                logger.warning(f"Could not load SFT adapter: {e}. Continuing with base only.")

        # Step 2: Load GRPO adapter on top of SFT-merged base
        grpo_id = config["grpo_adapter"]
        if grpo_id:
            logger.info(f"Loading GRPO adapter: {grpo_id}")
            try:
                model = PeftModel.from_pretrained(model, grpo_id, token=hf_token)
                logger.info("GRPO adapter loaded")
            except Exception as e:
                logger.warning(f"Could not load GRPO adapter: {e}. Using SFT-only model.")

        model.eval()
        _cache[model_key] = (model, tokenizer)
        return model, tokenizer


def make_local_generate_fn(model_key: str = "gemma") -> Callable:
    """Create a generate_fn that uses the specified model.

    Returns a function compatible with run_agent(generate_fn, ...).
    """
    model, tokenizer = load_model(model_key)
    config = MODEL_CONFIGS[model_key]
    chat_style = config["chat_style"]

    def generate(messages: list[dict[str, str]]) -> str:
        """Generate a response from the fine-tuned model."""
        import torch

        # Build the chat prompt
        chat = []
        for msg in messages:
            role = msg["role"]
            if chat_style == "gemma":
                # Gemma uses user/model roles; prepend system as user
                if role == "system":
                    chat.append({"role": "user", "content": msg["content"]})
                    chat.append({"role": "model", "content": "Understood. I will follow these instructions."})
                elif role == "assistant":
                    chat.append({"role": "model", "content": msg["content"]})
                else:
                    chat.append({"role": role, "content": msg["content"]})
            else:
                # ChatML / standard — pass through directly
                chat.append({"role": role, "content": msg["content"]})

        prompt = tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True
        )

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.1,
                do_sample=True,
                top_p=0.9,
            )

        # Decode only the new tokens
        new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        response = tokenizer.decode(new_tokens, skip_special_tokens=True)
        return response.strip()

    return generate


def make_base_generate_fn(model_key: str = "gemma") -> Callable:
    """Generate using the *base model* with adapter disabled.

    Uses the already-loaded PeftModel with ``model.disable_adapter()``
    so we get a text-only baseline on the exact same weights — no
    second download, no extra VRAM.

    The prompt asks for ``FINAL ANSWER: <name>`` so we can regex-parse
    the response the same way the frontier baselines work.
    """
    model, tokenizer = load_model(model_key)
    config = MODEL_CONFIGS[model_key]
    chat_style = config["chat_style"]

    def generate_base(narrative: str, question: str) -> tuple[str, str]:
        """Return (full_response, parsed_answer)."""
        import re
        import torch

        # Text-only prompt — no tools, no ontology, just the narrative
        system = (
            "You are an expert in Italian inheritance law. "
            "Read the family narrative carefully, then answer the question. "
            "End your response with exactly: FINAL ANSWER: <full name>"
        )
        user = f"{narrative}\n\n{question}"

        # Build chat
        chat = []
        if chat_style == "gemma":
            chat.append({"role": "user", "content": system})
            chat.append({"role": "model", "content": "Understood. I will follow these instructions."})
            chat.append({"role": "user", "content": user})
        else:
            chat.append({"role": "system", "content": system})
            chat.append({"role": "user", "content": user})

        prompt = tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True
        )
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        with torch.no_grad(), model.disable_adapter():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,
                temperature=1.0,
            )

        new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

        # Parse FINAL ANSWER
        match = re.search(r"FINAL ANSWER:\s*(.+?)(?:\n|$)", response, re.IGNORECASE)
        answer = match.group(1).strip().rstrip(".") if match else response.split("\n")[-1].strip()

        return response, answer

    return generate_base