feat: add OpenBMB Core Router (MiniCPM4.1-8B) and wire it to Gradio UI

agent/brain.py

@@ -0,0 +1,104 @@
+import os
+import sys
+import types
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# --- HACK TO BYPASS FLASH ATTENTION COMPILATION ---
+# MiniCPM4.1-8B strictly checks if flash_attn is installed, but compiling
+# it takes 20+ minutes on a cold HF Space. We mock the module exactly like
+# we did for the vision model so it passes the internal import check.
+if "flash_attn" not in sys.modules:
+    import importlib.machinery
+    _mock = types.ModuleType("flash_attn")
+    _mock.__spec__ = importlib.machinery.ModuleSpec("flash_attn", None)
+    _mock.__version__ = "2.6.0"
+    sys.modules["flash_attn"] = _mock
+
+# ---------------------------------------------------------------------------
+# Global Initialization — Core Boss (OpenBMB MiniCPM4.1-8B)
+# We load the text routing model into CPU memory when the space starts.
+# We do NOT call .cuda() here — ZeroGPU handles that inside @spaces.GPU.
+# ---------------------------------------------------------------------------
+print("Initializing Core Router (OpenBMB MiniCPM4.1-8B)... This may take a moment.")
+
+# Using the exact 8B model specified in our architecture
+model_id = "openbmb/MiniCPM4.1-8B" 
+hf_token = os.environ.get("HF_TOKEN")
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id, 
+        trust_remote_code=True, 
+        token=hf_token
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        attn_implementation="sdpa",  # Use built-in PyTorch attention, not flash_attn
+        token=hf_token
+    )
+except Exception as e:
+    print(f"Error loading Core Router: {e}")
+    tokenizer = None
+    model = None
+
+@spaces.GPU
+def _core_router_inference(prompt):
+    """
+    Executes the LLM logic on the ZeroGPU.
+    """
+    if model is None or tokenizer is None:
+        return "Error: Core Router Model failed to load."
+        
+    model.to("cuda")
+    model.eval()
+    
+    # Format the prompt using the chat template if available, else plain text
+    messages = [{"role": "user", "content": prompt}]
+    
+    # MiniCPM typically uses a specific chat template or plain tokenization
+    try:
+        # If the tokenizer supports apply_chat_template
+        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to("cuda")
+    except Exception:
+        # Fallback to standard tokenization
+        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+        inputs = inputs.input_ids
+        
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs, 
+            max_new_tokens=500,
+            do_sample=True,    # Required — without this, temperature/top_p crash
+            temperature=0.7,
+            top_p=0.9
+        )
+    
+    # Decode the newly generated tokens
+    input_length = inputs.shape[1] if hasattr(inputs, "shape") else len(inputs[0])
+    response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+    
+    return response
+
+def process_workflow(user_text, raw_vision_text=None):
+    """
+    The Main Flow Controller.
+    Orchestrates: Vision -> Parse -> Mellum -> Cohere -> VoxCPM2.
+    Additional parameters (clean_json, math_summary, translated_summary)
+    will be added as we build each specialist module.
+    """
+    
+    # Base prompt telling the model what it is
+    system_context = "You are the Core Router for the Family Bill Assistant, a helpful AI."
+    
+    # If we have extracted text from the receipt
+    if raw_vision_text:
+        prompt = f"{system_context}\n\nThe user uploaded a receipt. Here is the raw extracted text:\n{raw_vision_text}\n\nThe user says: {user_text}\n\nPlease respond helpfully."
+    else:
+        prompt = f"{system_context}\n\nThe user says: {user_text}\n\nPlease respond to their question."
+        
+    return _core_router_inference(prompt)

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from ui.layout import create_ui
 from tools.vision import process_receipt_image
+from agent.brain import process_workflow
 
 # Load the custom CSS for the "Off-Brand" Badge
 try:
@@ -20,23 +21,27 @@ with demo:
     # Initialize the UI layout from the ui folder
     image_input, audio_input, submit_btn, chatbot, msg_input = create_ui()
     
-    # Bind the submit button to the vision function
-    def handle_analyze(image):
-        if not image:
-            return [{"role": "assistant", "content": "Please upload an image first!"}]
+    # Bind the submit button to the Core Boss workflow
+    def handle_analyze(image, user_msg, history):
+        if not user_msg:
+            user_msg = "Please analyze this bill."
+            
+        # Step 1: If an image is provided, extract raw text
+        raw_text = None
+        if image:
+            raw_text = process_receipt_image(image)
+            
+        # Step 2: Route everything to the Core Boss
+        bot_response = process_workflow(user_text=user_msg, raw_vision_text=raw_text)
         
-        # Trigger the ZeroGPU
-        result = process_receipt_image(image)
-        
-        # Return as a chat message history
-        return [
-            {"role": "user", "content": "Here is my bill. Can you extract the text?"},
-            {"role": "assistant", "content": str(result)}
-        ]
+        # Step 3: Append to chat history
+        history.append({"role": "user", "content": user_msg})
+        history.append({"role": "assistant", "content": str(bot_response)})
+        return history
         
     submit_btn.click(
         fn=handle_analyze,
-        inputs=[image_input],
+        inputs=[image_input, msg_input, chatbot],
         outputs=[chatbot]
     )