agent/brain.py

import os
import sys
import spaces
import torch
from unittest.mock import patch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.dynamic_module_utils import get_imports as _original_get_imports

# Capture original BEFORE any patching happens.
# If imported INSIDE the patched function it recurses into itself infinitely.
def _strip_flash_attn(filename, *args, **kwargs):
    """Strips flash_attn from model import list so FA2 is never enabled."""
    return [i for i in _original_get_imports(filename, *args, **kwargs) if i != "flash_attn"]

# ---------------------------------------------------------------------------
# BLOCK flash_attn at RUNTIME level
#
# The get_imports patch above strips flash_attn from the requirements list
# (so transformers can load the module files). But MiniCPM4.1-8B's own
# modeling_minicpm.py also does a RUNTIME check at model init:
#
#   try: import flash_attn; HAS_FLASH_ATTN = True
#   except ImportError: HAS_FLASH_ATTN = False
#
# HuggingFace ZeroGPU machines have flash_attn pre-installed, so this
# runtime check would succeed and force FA2 on. Setting sys.modules entry
# to None is Python's standard way to permanently block a module import —
# any subsequent "import flash_attn" will raise ImportError.
# ---------------------------------------------------------------------------
sys.modules["flash_attn"] = None  # None entry = ImportError on any import attempt

# ---------------------------------------------------------------------------
# COMPATIBILITY PATCH — restore is_torch_fx_available
#
# MiniCPM4.1-8B's custom modeling_minicpm.py imports is_torch_fx_available
# from transformers.utils.import_utils. This function was removed in
# transformers>=4.47 but is still referenced in the model's remote code.
# We restore it BEFORE anything else runs.
# ---------------------------------------------------------------------------
try:
    import transformers.utils.import_utils as _tu
    if not hasattr(_tu, "is_torch_fx_available"):
        _tu.is_torch_fx_available = lambda: False
except Exception:
    pass

# ---------------------------------------------------------------------------
# LAZY LOADING — Core Boss (OpenBMB MiniCPM4.1-8B)
#
# WHY LAZY?
# MiniCPM4.1-8B is ~16GB of weights. Loading it at startup exhausts the
# HuggingFace Space startup timeout. By loading it on the FIRST user
# request instead, the Space starts instantly and users see the UI
# immediately. The first request takes longer (model download), but all
# subsequent requests are instant.
# ---------------------------------------------------------------------------

model_id = "openbmb/MiniCPM4.1-8B"
hf_token = os.environ.get("HF_TOKEN")

_tokenizer = None
_model = None


def _patch_flash_attn_imports(filename, *args, **kwargs):
    """
    Documented community fix: patch transformers' get_imports to strip
    flash_attn from the model's required imports before from_pretrained runs.
    This prevents MiniCPM4.1-8B's modeling_minicpm.py from forcing FA2
    regardless of attn_implementation setting.
    Source: https://huggingface.co/openbmb/MiniCPM4.1-8B/discussions
    """
    from transformers.dynamic_module_utils import get_imports as _orig_get_imports
    imports = _orig_get_imports(filename, *args, **kwargs)
    return [imp for imp in imports if imp != "flash_attn"]


def _ensure_loaded():
    """Thread-safe lazy loader. Does nothing after first successful load."""
    global _tokenizer, _model

    if _model is not None:
        return True

    print("[Core Router] First request received — loading MiniCPM4.1-8B now...")
    try:
        _tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True,
            token=hf_token
        )

        # Patch get_imports to remove flash_attn from the dependency list.
        # This is the ONLY reliable way to prevent FA2 from being forced
        # by the model's own custom code, regardless of attn_implementation.
        with patch(
            "transformers.dynamic_module_utils.get_imports",
            side_effect=_strip_flash_attn
        ):
            _model = AutoModelForCausalLM.from_pretrained(
                model_id,
                trust_remote_code=True,
                torch_dtype="auto",
                low_cpu_mem_usage=True,
                attn_implementation="eager",
                token=hf_token
            )

        # Removed the forced ChatML template so the model uses its native tags

        print("[Core Router] MiniCPM4.1-8B loaded successfully!")
        return True
    except Exception as e:
        print(f"[Core Router] Load failed: {e}")
        _model = None
        _tokenizer = None
        return False


@spaces.GPU
def _core_router_inference(prompt: str) -> str:
    """
    Runs MiniCPM4.1-8B inference on the ZeroGPU.
    Lazy-loads the model on the first call.
    """
    if not _ensure_loaded():
        return "Error: Core Router model could not be loaded. Please try again."

    _model.to("cuda")
    _model.eval()

    try:
        # --- SMOKE TEST OVERRIDE ---
        prompt = "Hello, what is 2+2? Please answer clearly."
        print("\n\n" + "="*50)
        print("=== CHRISTIAN'S SMOKE TEST ===")
        print(f"1. tokenizer.chat_template:\n{_tokenizer.chat_template}\n")
        
        # 1. Use the tokenizer's native chat template
        messages = [{"role": "user", "content": prompt}]
        prompt_str = _tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
        print(f"2. Decoded prompt before generation:\n{prompt_str}\n")
        
        # 2. Tokenize and move to GPU
        inputs = _tokenizer(prompt_str, return_tensors="pt").to("cuda")
        print(f"3. Input length: {inputs.input_ids.shape}\n")
        
        # 3. Deterministic Generation (as requested by Christian)
        print("4. Starting Deterministic Generation...")
        outputs = _model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=False,
            eos_token_id=_tokenizer.eos_token_id,
            pad_token_id=_tokenizer.pad_token_id if _tokenizer.pad_token_id is not None else _tokenizer.eos_token_id,
            repetition_penalty=1.05,
            use_cache=False
        )
        
        # 4. Slice off the prompt tokens so we only decode the generated answer
        generated_ids = outputs[0][inputs.input_ids.shape[1]:]
        print(f"5. Raw generated IDs:\n{generated_ids.tolist()}\n")
        
        # 5. Decode back to text
        response = _tokenizer.decode(generated_ids, skip_special_tokens=True)
        print(f"6. Decoded continuation only:\n{response}")
        print("="*50 + "\n\n")
        
        return response.strip()
    except Exception as e:
        return f"Error during Core Router generation: {e}"


def process_workflow(user_text: str, raw_vision_text: str = None) -> str:
    """
    The Main Flow Controller.
    Orchestrates: Vision -> Parse -> Mellum -> Cohere -> VoxCPM2.
    """
    system_context = (
        "You are the Core Router for the Family Bill Assistant. "
        "You are a helpful, friendly AI that helps families understand their bills. "
        "You can analyze receipts, split bills, and answer financial questions."
    )

    if raw_vision_text:
        prompt = (
            f"{system_context}\n\n"
            f"The user uploaded a receipt. Here is the extracted text:\n"
            f"{raw_vision_text}\n\n"
            f"The user says: {user_text}\n\n"
            f"Please analyze the receipt and respond helpfully."
        )
    else:
        prompt = (
            f"{system_context}\n\n"
            f"The user says: {user_text}\n\n"
            f"Please respond to their question."
        )

    return _core_router_inference(prompt)