Delete shared_models.py

shared_models.py

@@ -1,285 +0,0 @@
-# shared_models.py
-"""
-Shared model manager for Mimir agents.
-Uses Llama-3.2-3B-Instruct with transformers for all agents.
-"""
-import torch
-import threading
-import logging
-import os
-from typing import Optional, List
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-logger = logging.getLogger(__name__)
-
-# ZeroGPU support
-try:
-    import spaces
-    HF_SPACES_AVAILABLE = True
-    logger.info("✅ ZeroGPU (spaces) available")
-except ImportError:
-    HF_SPACES_AVAILABLE = False
-    class DummySpaces:
-        @staticmethod
-        def GPU(duration=90):
-            def decorator(func):
-                return func
-            return decorator
-    spaces = DummySpaces()
-    logger.warning("⚠️ ZeroGPU not available - running without GPU allocation")
-
-HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
-# Model configuration
-LLAMA_MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
-
-
-class LlamaSharedAgent:
-    """
-    Singleton agent using Llama-3.2-3B-Instruct for all Mimir operations.
-    Thread-safe with ZeroGPU allocation management.
-    
-    Used by:
-    - ToolDecisionAgent
-    - PromptRoutingAgents (all 4 agents)
-    - ThinkingAgents (all reasoning agents)
-    - ResponseAgent
-    """
-    
-    _instance = None
-    _lock = threading.Lock()
-    
-    def __new__(cls):
-        """Ensure only one instance exists (singleton pattern)"""
-        if cls._instance is None:
-            with cls._lock:
-                if cls._instance is None:
-                    cls._instance = super().__new__(cls)
-        return cls._instance
-    
-    def __init__(self):
-        """Initialize only once"""
-        if hasattr(self, '_initialized'):
-            return
-        
-        self.pipe = None
-        self.tokenizer = None
-        self.model = None
-        self.model_loaded = False
-        self._initialized = True
-        logger.info("LlamaSharedAgent instance created (singleton)")
-    
-    @spaces.GPU(duration=120)
-    def _ensure_loaded(self):
-        """
-        Load model with GPU allocation (ZeroGPU).
-        Only ONE @spaces.GPU decorator for Llama across entire app!
-        """
-        if self.model_loaded:
-            logger.info("✅ Llama-3.2-3B already loaded, reusing existing instance")
-            return
-        
-        logger.info("="*60)
-        logger.info("LOADING SHARED LLAMA-3.2-3B-INSTRUCT")
-        logger.info("="*60)
-        
-        try:
-            # 4-bit quantization config for memory efficiency
-            quantization_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_use_double_quant=True,
-            )
-            
-            logger.info(f"Loading model: {LLAMA_MODEL_ID}")
-            logger.info("Configuration: 4-bit NF4 quantization")
-            
-            # Load tokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                LLAMA_MODEL_ID,
-                token=HF_TOKEN,
-                trust_remote_code=True,
-            )
-            
-            # Load model with quantization
-            self.model = AutoModelForCausalLM.from_pretrained(
-                LLAMA_MODEL_ID,
-                quantization_config=quantization_config,
-                device_map="auto",
-                token=HF_TOKEN,
-                trust_remote_code=True,
-                torch_dtype=torch.bfloat16,
-            )
-            
-            # Create pipeline
-            self.pipe = pipeline(
-                "text-generation",
-                model=self.model,
-                tokenizer=self.tokenizer,
-                torch_dtype=torch.bfloat16,
-                device_map="auto",
-            )
-            
-            self.model_loaded = True
-            
-            logger.info("="*60)
-            logger.info("✅ SHARED LLAMA-3.2-3B LOADED SUCCESSFULLY")
-            logger.info(f"  Model: {LLAMA_MODEL_ID}")
-            logger.info(f"  Quantization: 4-bit NF4")
-            logger.info(f"  Memory: ~1GB VRAM (vs 3.3GB GGUF)")
-            logger.info("  Context: 128K tokens")
-            logger.info("  This model will be reused by:")
-            logger.info("  - ToolDecisionAgent")
-            logger.info("  - PromptRoutingAgents (all 4 agents)")
-            logger.info("  - ThinkingAgents (all reasoning)")
-            logger.info("  - ResponseAgent (final responses)")
-            logger.info("="*60)
-            
-        except Exception as e:
-            logger.error(f"Failed to load Llama-3.2-3B: {e}")
-            raise
-    
-    def generate(
-        self, 
-        system_prompt: str,
-        user_message: str,
-        max_tokens: int = 100,
-        temperature: float = 0.7,
-        stop_sequences: Optional[List[str]] = None
-    ) -> str:
-        """
-        Generate response using shared Llama-3.2-3B model.
-        
-        Args:
-            system_prompt: System instruction
-            user_message: User query
-            max_tokens: Max tokens to generate
-            temperature: Sampling temperature
-            stop_sequences: Optional list of stop sequences (not used with pipeline)
-            
-        Returns:
-            Generated text
-        """
-        # Ensure model is loaded (triggers @spaces.GPU only once)
-        self._ensure_loaded()
-        
-        # Format messages using Llama 3.2 chat template (handled automatically)
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_message},
-        ]
-        
-        try:
-            # Generate using pipeline
-            outputs = self.pipe(
-                messages,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                do_sample=True,
-                top_p=0.9,
-                top_k=40,
-                repetition_penalty=1.15,
-            )
-            
-            # Extract generated text (pipeline returns full conversation)
-            result = outputs[0]["generated_text"][-1]["content"]
-            
-            return result.strip()
-            
-        except Exception as e:
-            logger.error(f"Generation error: {e}")
-            return ""
-    
-    def generate_streaming(
-        self,
-        system_prompt: str,
-        user_message: str,
-        max_tokens: int = 512,
-        temperature: float = 0.7,
-    ):
-        """
-        Generate response with streaming (for ResponseAgent).
-        
-        Yields:
-            str: Generated text chunks
-        """
-        self._ensure_loaded()
-        
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_message},
-        ]
-        
-        try:
-            # Use TextIteratorStreamer for streaming
-            from transformers import TextIteratorStreamer
-            from threading import Thread
-            
-            # Apply chat template
-            input_ids = self.tokenizer.apply_chat_template(
-                messages,
-                add_generation_prompt=True,
-                return_tensors="pt"
-            ).to(self.model.device)
-            
-            streamer = TextIteratorStreamer(
-                self.tokenizer,
-                skip_prompt=True,
-                skip_special_tokens=True
-            )
-            
-            generation_kwargs = dict(
-                input_ids=input_ids,
-                streamer=streamer,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                do_sample=True,
-                top_p=0.9,
-                top_k=40,
-                repetition_penalty=1.15,
-            )
-            
-            # Generate in separate thread
-            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
-            thread.start()
-            
-            # Yield generated text
-            for text in streamer:
-                yield text
-                
-        except Exception as e:
-            logger.error(f"Streaming generation error: {e}")
-            yield ""
-    
-    def get_model_info(self) -> dict:
-        """Get model information for diagnostics"""
-        return {
-            "status": "loaded" if self.model_loaded else "not_loaded",
-            "model_id": LLAMA_MODEL_ID,
-            "model_type": "llama-3.2-3b-instruct",
-            "quantization": "4-bit NF4",
-            "size_gb": 1.0,
-            "context_length": 128000,
-            "zerogpu_ready": True,
-            "transformers_pipeline": True,
-            "shared_instance": True,
-        }
-
-
-# Global singleton instance
-_shared_llama = None
-
-def get_shared_llama() -> LlamaSharedAgent:
-    """Get or create the shared Llama-3.2-3B agent instance"""
-    global _shared_llama
-    if _shared_llama is None:
-        _shared_llama = LlamaSharedAgent()
-    return _shared_llama
-
-
-# Backwards compatibility aliases
-Qwen3SharedAgent = LlamaSharedAgent
-MistralSharedAgent = LlamaSharedAgent
-get_shared_qwen3 = get_shared_llama
-get_shared_mistral = get_shared_llama