"""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: `` 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: " ) 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