Upload app.py

app.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
-renamed from app_real.py - Production-Ready Mamba Encoder Swarm Demo
-Combines real model functionality with rich UI and comprehensive error handling
+Enhanced Production-Ready Mamba Encoder Swarm Demo
+Integrates pretrained Mamba weights from HuggingFace with swarm architecture
 """
 
 import gradio as gr
@@ -14,7 +14,8 @@ import os
 import psutil
 from typing import Optional, Dict, Any, Tuple
 from datetime import datetime
-from transformers import AutoTokenizer, AutoConfig
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
+from huggingface_hub import snapshot_download, hf_hub_download
 
 # Setup comprehensive logging
 logging.basicConfig(
@@ -27,8 +28,106 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+class MambaWeightLoader:
+    """Dynamic loader for pretrained Mamba weights"""
+    
+    def __init__(self, model_name="state-spaces/mamba-130m"):
+        self.model_name = model_name
+        self.cache_dir = "/tmp/mamba_cache" if os.path.exists("/tmp") else "./mamba_cache"
+        self.model = None
+        self.tokenizer = None
+        self.config = None
+    
+    def download_and_load(self):
+        """Download and load Mamba weights in HuggingFace Spaces"""
+        try:
+            logger.info(f"🔄 Loading pretrained model: {self.model_name}")
+            
+            # Create cache directory
+            os.makedirs(self.cache_dir, exist_ok=True)
+            
+            # Load tokenizer (lightweight)
+            logger.info("📝 Loading tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name,
+                cache_dir=self.cache_dir,
+                trust_remote_code=True
+            )
+            
+            # Handle tokenizer padding
+            if self.tokenizer.pad_token is None:
+                if self.tokenizer.eos_token is not None:
+                    self.tokenizer.pad_token = self.tokenizer.eos_token
+                else:
+                    self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+            
+            # Load configuration
+            logger.info("⚙️ Loading model configuration...")
+            self.config = AutoConfig.from_pretrained(
+                self.model_name,
+                cache_dir=self.cache_dir,
+                trust_remote_code=True
+            )
+            
+            # Load model with optimizations for Spaces
+            logger.info("🧠 Loading model weights...")
+            
+            # Determine optimal dtype and device settings
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            dtype = torch.float16 if device.type == "cuda" else torch.float32
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                config=self.config,
+                cache_dir=self.cache_dir,
+                trust_remote_code=True,
+                torch_dtype=dtype,
+                device_map="auto" if torch.cuda.is_available() else None,
+                low_cpu_mem_usage=True
+            )
+            
+            # Move to device if not using device_map
+            if not torch.cuda.is_available():
+                self.model.to(device)
+            
+            self.model.eval()
+            
+            # Log model info
+            num_params = sum(p.numel() for p in self.model.parameters())
+            logger.info(f"✅ Model loaded successfully!")
+            logger.info(f"📊 Parameters: {num_params:,} ({num_params/1e6:.1f}M)")
+            logger.info(f"🔧 Device: {device}, dtype: {dtype}")
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Error loading pretrained model: {e}")
+            return False
+    
+    def get_model_info(self):
+        """Get model information"""
+        if self.model:
+            try:
+                num_params = sum(p.numel() for p in self.model.parameters())
+                device = next(self.model.parameters()).device
+                dtype = next(self.model.parameters()).dtype
+                
+                return {
+                    "name": self.model_name,
+                    "parameters": f"{num_params:,}",
+                    "parameters_millions": f"{num_params/1e6:.1f}M",
+                    "device": str(device),
+                    "dtype": str(dtype),
+                    "vocab_size": getattr(self.config, 'vocab_size', 'Unknown'),
+                    "hidden_size": getattr(self.config, 'd_model', getattr(self.config, 'hidden_size', 'Unknown'))
+                }
+            except Exception as e:
+                logger.error(f"Error getting model info: {e}")
+                return {"error": str(e)}
+        return None
+
 class MambaSwarmDemo:
-    """Production-ready Mamba Swarm Demo with fallback capabilities"""
+    """Enhanced Production-ready Mamba Swarm Demo with dynamic pretrained weight loading"""
     
     def __init__(self, model_path: str = "./", fallback_mode: bool = False):
         self.model = None
@@ -38,6 +137,8 @@ class MambaSwarmDemo:
         self.model_path = model_path
         self.fallback_mode = fallback_mode
         self.model_loaded = False
+        self.pretrained_loader = None
+        self.using_pretrained = False
         
         # Performance tracking
         self.stats = {
@@ -60,24 +161,23 @@ class MambaSwarmDemo:
         }
         
         self._initialize_model()
-        logger.info(f"Demo initialized - Model loaded: {self.model_loaded}, Fallback mode: {self.fallback_mode}")
+        logger.info(f"Demo initialized - Model loaded: {self.model_loaded}, Using pretrained: {self.using_pretrained}, Fallback mode: {self.fallback_mode}")
     
     def _initialize_model(self):
-        """Initialize model with comprehensive error handling and fallback"""
+        """Initialize model with pretrained weights or fallback"""
         try:
-            logger.info("Attempting to load Mamba Swarm model...")
+            logger.info("🚀 Attempting to load model with priority: Pretrained -> Custom -> Fallback")
             
-            # Check if model files exist
-            config_path = os.path.join(self.model_path, "config.json")
-            if not os.path.exists(config_path) and not self.fallback_mode:
-                logger.warning(f"Config file not found at {config_path}, enabling fallback mode")
-                self.fallback_mode = True
+            # Try to load pretrained model first (highest priority)
+            success = self._load_pretrained_model()
             
-            if not self.fallback_mode:
-                # Try to load real model
-                self._load_real_model()
-            else:
-                # Initialize in fallback mode
+            if not success:
+                logger.info("Pretrained loading failed, trying custom swarm model...")
+                success = self._load_custom_swarm_model()
+            
+            if not success:
+                logger.info("All model loading attempts failed, enabling fallback mode")
+                self.fallback_mode = True
                 self._initialize_fallback_mode()
                 
         except Exception as e:
@@ -86,132 +186,136 @@ class MambaSwarmDemo:
             self.fallback_mode = True
             self._initialize_fallback_mode()
     
-    def _load_real_model(self):
-        """Load the actual Mamba Swarm model"""
+    def _load_pretrained_model(self):
+        """Load pretrained Mamba model from HuggingFace with automatic model selection"""
+        try:
+            # Choose model based on available resources
+            MODEL_OPTIONS = {
+                "small": "state-spaces/mamba-130m",      # ~500MB
+                "medium": "state-spaces/mamba-790m",     # ~3GB  
+                "large": "state-spaces/mamba-1.4b",      # ~5GB
+                "xl": "state-spaces/mamba-2.8b",         # ~10GB
+            }
+            
+            # Auto-select model based on available memory
+            memory_gb = psutil.virtual_memory().total / (1024**3)
+            if memory_gb >= 32 and torch.cuda.is_available():
+                selected_model = MODEL_OPTIONS["xl"]
+            elif memory_gb >= 16 and torch.cuda.is_available():
+                selected_model = MODEL_OPTIONS["large"]
+            elif memory_gb >= 8:
+                selected_model = MODEL_OPTIONS["medium"]
+            else:
+                selected_model = MODEL_OPTIONS["small"]
+            
+            logger.info(f"🎯 Auto-selected model: {selected_model} (Available memory: {memory_gb:.1f}GB)")
+            
+            # Initialize loader
+            self.pretrained_loader = MambaWeightLoader(selected_model)
+            
+            # Download and load
+            if self.pretrained_loader.download_and_load():
+                self.model = self.pretrained_loader.model
+                self.tokenizer = self.pretrained_loader.tokenizer
+                self.config = self.pretrained_loader.config
+                self.model_loaded = True
+                self.using_pretrained = True
+                
+                logger.info("✅ Pretrained model loaded successfully!")
+                return True
+            else:
+                logger.warning("❌ Pretrained model loading failed")
+                return False
+                
+        except Exception as e:
+            logger.error(f"Pretrained model loading error: {e}")
+            return False
+    
+    def _load_custom_swarm_model(self):
+        """Try to load custom swarm model implementation"""
         try:
-            # Try multiple import paths for the model
+            logger.info("Attempting to load custom Mamba Swarm model...")
+            
+            # Try multiple import paths for the custom model
             model_class = None
             
-            # Try importing from different locations
             try:
                 from modeling_mamba_swarm import MambaSwarmForCausalLM
                 model_class = MambaSwarmForCausalLM
-                logger.info("Loaded MambaSwarmForCausalLM from modeling_mamba_swarm")
+                logger.info("Found MambaSwarmForCausalLM")
             except ImportError:
                 try:
-                    from upload_to_hf import MambaSwarmForCausalLM
-                    model_class = MambaSwarmForCausalLM
-                    logger.info("Loaded MambaSwarmForCausalLM from upload_to_hf")
+                    from core.mamba_swarm_integration import MambaEncoderSwarmModel
+                    model_class = MambaEncoderSwarmModel
+                    logger.info("Found MambaEncoderSwarmModel")
                 except ImportError:
                     try:
-                        from core.mamba_swarm_integration import MambaEncoderSwarmModel
-                        model_class = MambaEncoderSwarmModel
-                        logger.info("Loaded MambaEncoderSwarmModel from core.mamba_swarm_integration")
+                        from system.mambaSwarm import UnifiedMambaSwarm
+                        # Use the unified swarm in native mode
+                        swarm = UnifiedMambaSwarm(use_pretrained=False)
+                        if hasattr(swarm, 'native_swarm_model') and swarm.native_swarm_model:
+                            self.model = swarm.native_swarm_model
+                            self.model_loaded = True
+                            logger.info("Loaded native swarm model from UnifiedMambaSwarm")
+                            return True
+                        else:
+                            raise ImportError("No native swarm model available")
                     except ImportError:
-                        try:
-                            from system.mambaSwarm import UnifiedMambaSwarm
-                            # Use the unified swarm in native mode
-                            swarm = UnifiedMambaSwarm(use_pretrained=False)
-                            if hasattr(swarm, 'native_swarm_model') and swarm.native_swarm_model:
-                                self.model = swarm.native_swarm_model
-                                self.model_loaded = True
-                                logger.info("Loaded native swarm model from UnifiedMambaSwarm")
-                                return
-                            else:
-                                raise ImportError("No native swarm model available")
-                        except ImportError as e:
-                            logger.error(f"All model imports failed: {e}")
-                            raise ImportError("No compatible Mamba Swarm model found")
+                        logger.warning("No custom swarm model found")
+                        return False
             
             if model_class is None:
-                raise ImportError("No model class available")
+                return False
             
-            # Load configuration
+            # Create configuration for custom model
             try:
-                self.config = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True)
-                logger.info(f"Loaded config: {self.config.__class__.__name__}")
-            except Exception as e:
-                logger.warning(f"Could not load config from {self.model_path}: {e}")
-                # Create a default config using our MambaSwarmConfig
+                from modeling_mamba_swarm import MambaSwarmConfig
+                self.config = MambaSwarmConfig(
+                    num_encoders=8,
+                    max_mamba_encoders=100,
+                    d_model=768,
+                    vocab_size=50257,
+                    max_sequence_length=2048
+                )
+            except ImportError:
+                # Fallback config
                 try:
-                    from modeling_mamba_swarm import MambaSwarmConfig
-                    self.config = MambaSwarmConfig(
-                        num_encoders=8,
-                        max_mamba_encoders=100,
-                        d_model=768,
-                        vocab_size=50257,
-                        max_sequence_length=2048
-                    )
-                    logger.info("Using default MambaSwarmConfig")
-                except ImportError:
-                    # Final fallback to basic config
                     from core.config import MambaConfig
                     self.config = MambaConfig()
-                    # Add swarm-specific attributes
                     self.config.num_encoders = 8
                     self.config.max_mamba_encoders = 100
-                    self.config.max_sequence_length = 2048
-                    logger.info("Using default MambaConfig with swarm attributes")
+                except ImportError:
+                    # Create minimal config
+                    self.config = type('Config', (), {
+                        'num_encoders': 8,
+                        'max_mamba_encoders': 100,
+                        'd_model': 768,
+                        'vocab_size': 50257,
+                        'max_sequence_length': 2048
+                    })()
             
-            # Load tokenizer
-            try:
-                self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
-                if self.tokenizer.pad_token is None:
-                    self.tokenizer.pad_token = self.tokenizer.eos_token
-                logger.info("Tokenizer loaded successfully")
-            except Exception as e:
-                logger.warning(f"Could not load tokenizer: {e}")
-                # Use a simple fallback tokenizer
-                from transformers import GPT2Tokenizer
-                self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-                if self.tokenizer.pad_token is None:
-                    self.tokenizer.pad_token = self.tokenizer.eos_token
-                logger.info("Using fallback GPT2 tokenizer")
-            
-            # Load model with memory optimization
-            dtype = torch.float16 if self.device.type == "cuda" else torch.float32
-            
-            if model_class == MambaEncoderSwarmModel:
-                # Native integration model - create with MambaConfig
-                from core.config import MambaConfig
-                if not hasattr(self, 'config') or not isinstance(self.config, MambaConfig):
-                    mamba_config = MambaConfig(
-                        d_model=getattr(self.config, 'd_model', 768),
-                        vocab_size=getattr(self.config, 'vocab_size', 50257),
-                        n_layers=8,
-                        d_state=16,
-                        d_conv=4,
-                        bias=False
-                    )
-                    self.model = model_class(mamba_config, num_encoders=getattr(self.config, 'num_encoders', 8))
-                else:
-                    self.model = model_class(self.config, num_encoders=getattr(self.config, 'num_encoders', 8))
+            # Initialize custom model
+            if model_class.__name__ == 'MambaEncoderSwarmModel':
+                self.model = model_class(self.config, num_encoders=8)
             else:
-                # HuggingFace-style model or our new MambaSwarmForCausalLM
-                if hasattr(model_class, 'from_pretrained') and os.path.exists(self.model_path):
-                    self.model = model_class.from_pretrained(
-                        self.model_path,
-                        config=self.config,
-                        torch_dtype=dtype,
-                        trust_remote_code=True,
-                        low_cpu_mem_usage=True
-                    )
-                else:
-                    # Create with config only
-                    self.model = model_class(self.config)
+                self.model = model_class(self.config)
+            
+            # Create tokenizer
+            from transformers import GPT2Tokenizer
+            self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
             
             self.model.to(self.device)
             self.model.eval()
             self.model_loaded = True
             
-            # Log model info
-            num_params = sum(p.numel() for p in self.model.parameters())
-            logger.info(f"Model loaded successfully on {self.device}")
-            logger.info(f"Model parameters: {num_params:,} ({num_params/1e6:.1f}M)")
+            logger.info("✅ Custom swarm model loaded successfully!")
+            return True
             
         except Exception as e:
-            logger.error(f"Real model loading failed: {e}")
-            raise
+            logger.error(f"Custom model loading error: {e}")
+            return False
     
     def _initialize_fallback_mode(self):
         """Initialize fallback/simulation mode"""
@@ -246,7 +350,6 @@ class MambaSwarmDemo:
                 self.eos_token = "[EOS]"
             
             def encode(self, text, return_tensors=None):
-                # Simple word-based tokenization for simulation
                 tokens = text.split()
                 token_ids = [hash(token) % 1000 for token in tokens]
                 if return_tensors == "pt":
@@ -254,7 +357,6 @@ class MambaSwarmDemo:
                 return token_ids
             
             def decode(self, token_ids, skip_special_tokens=True):
-                # Mock decoding
                 return f"Generated response for {len(token_ids)} tokens"
         
         self.tokenizer = MockTokenizer()
@@ -310,7 +412,7 @@ class MambaSwarmDemo:
         available_encoders = list(range(start, min(end + 1, 101)))
         
         # Select encoders based on prompt complexity and domain
-        prompt_complexity = min(len(prompt.split()) / 10, 3.0)  # Complexity factor
+        prompt_complexity = min(len(prompt.split()) / 10, 3.0)
         optimal_count = min(max(int(num_encoders * (1 + prompt_complexity)), 3), 25)
         
         if len(available_encoders) >= optimal_count:
@@ -333,171 +435,9 @@ class MambaSwarmDemo:
             'total_active': len(selected_encoders)
         }
     
-    def _simulate_generation(self, prompt: str, routing_info: Dict, max_length: int) -> str:
-        """Generate sophisticated simulated responses based on domain"""
-        domain = routing_info['detected_domain']
-        
-        domain_responses = {
-            'medical': f"""Based on medical literature and current research, regarding "{prompt[:50]}...":
-
-This condition/topic involves multiple factors including genetic predisposition, environmental influences, and lifestyle factors. Key considerations include:
-
-• Proper medical evaluation is essential
-• Individual symptoms may vary significantly  
-• Treatment approaches should be personalized
-• Regular monitoring is typically recommended
-
-**Important**: This information is for educational purposes only. Please consult with qualified healthcare professionals for personalized medical advice and treatment recommendations.""",
-            
-            'legal': f"""From a legal perspective on "{prompt[:50]}...":
-
-The legal framework surrounding this matter involves several key considerations:
-
-• Jurisdictional requirements and applicable statutes
-• Precedent cases and regulatory guidelines
-• Compliance obligations and reporting requirements
-• Risk assessment and mitigation strategies
-
-**Disclaimer**: This information is for general informational purposes only and does not constitute legal advice. Consult with qualified legal professionals for specific legal matters.""",
-            
-            'code': f"""Here's a comprehensive solution for "{prompt[:50]}...":
-
-```python
-def optimized_solution(input_data):
-    \"\"\"
-    Efficient implementation with error handling
-    Time complexity: O(n log n)
-    Space complexity: O(n)
-    \"\"\"
-    try:
-        # Input validation
-        if not input_data:
-            raise ValueError("Input data cannot be empty")
-        
-        # Core algorithm implementation
-        result = process_data(input_data)
-        
-        # Additional optimizations
-        result = optimize_output(result)
-        
-        return result
-    
-    except Exception as e:
-        logger.error(f"Processing error: {{e}}")
-        return None
-
-def process_data(data):
-    # Implementation details here
-    return processed_data
-
-def optimize_output(data):
-    # Performance optimizations
-    return optimized_data
-```
-
-**Key Features:**
-• Error handling and input validation
-• Optimized performance characteristics
-• Comprehensive documentation
-• Production-ready implementation""",
-            
-            'science': f"""Scientific Analysis of "{prompt[:50]}...":
-
-Based on current scientific understanding and peer-reviewed research:
-
-**Theoretical Framework:**
-The underlying principles involve complex interactions between multiple variables, governed by established scientific laws and emerging theories.
-
-**Methodology:**
-• Systematic observation and data collection
-• Controlled experimental design
-• Statistical analysis and validation
-• Peer review and reproducibility testing
-
-**Current Research:**
-Recent studies indicate significant progress in understanding the mechanisms involved, with several promising avenues for future investigation.
-
-**Implications:**
-These findings have broad applications across multiple disciplines and may lead to significant advances in the field.""",
-            
-            'creative': f"""**{prompt[:30]}...**
-
-The story unfolds in a world where imagination meets reality, where every character carries the weight of their dreams and the burden of their choices.
-
-*Chapter 1: The Beginning*
-
-In the quiet moments before dawn, when the world holds its breath between night and day, our tale begins. The protagonist stands at the threshold of an adventure that will challenge everything they thought they knew about themselves and the world around them.
-
-The narrative weaves through layers of meaning, exploring themes of identity, purpose, and the delicate balance between hope and reality. Each scene is crafted with careful attention to emotional resonance and character development.
-
-*As the story progresses, we discover that the true journey is not external, but internal—a transformation of the soul that mirrors the changing landscape of the world itself.*
-
-**Themes Explored:**
-• Personal growth and self-discovery
-• The power of resilience and determination
-• The complexity of human relationships
-• The intersection of dreams and reality""",
-            
-            'business': f"""**Strategic Analysis: {prompt[:50]}...**
-
-**Executive Summary:**
-This comprehensive analysis examines the strategic implications and market opportunities related to the identified business challenge.
-
-**Market Assessment:**
-• Current market size and growth projections
-• Competitive landscape analysis
-• Key trends and disruption factors
-• Customer segment identification
-
-**Strategic Recommendations:**
-1. **Short-term actions** (0-6 months)
-   - Immediate market positioning
-   - Resource allocation optimization
-   - Risk mitigation strategies
-
-2. **Medium-term initiatives** (6-18 months)
-   - Strategic partnerships and alliances
-   - Product/service development
-   - Market expansion opportunities
-
-3. **Long-term vision** (18+ months)
-   - Innovation and R&D investment
-   - Scalability and sustainability
-   - Market leadership positioning
-
-**Financial Projections:**
-Based on conservative estimates, implementation of these strategies could result in significant ROI and market share growth.""",
-            
-            'general': f"""**Comprehensive Response to: "{prompt[:50]}..."**
-
-Thank you for your inquiry. Based on available knowledge and expertise from {routing_info['total_active']} specialized domains, here's a comprehensive analysis:
-
-**Key Points:**
-• The topic involves multiple interconnected factors that require careful consideration
-• Current understanding is based on established principles and ongoing research
-• Practical applications vary depending on specific context and requirements
-• Best practices emphasize a balanced, evidence-based approach
-
-**Detailed Analysis:**
-The subject matter encompasses several important dimensions that merit thorough examination. Each aspect contributes to a deeper understanding of the overall concept and its implications.
-
-**Practical Considerations:**
-Implementation requires careful planning, adequate resources, and ongoing monitoring to ensure optimal outcomes. Success factors include stakeholder engagement, clear communication, and adaptive management strategies.
-
-**Conclusion:**
-This analysis provides a foundation for informed decision-making while acknowledging the complexity and nuanced nature of the topic."""
-        }
-        
-        return domain_responses.get(domain, domain_responses['general'])
-    
     def generate_text(self, prompt: str, max_length: int = 100, temperature: float = 0.7, 
                      top_p: float = 0.9, num_encoders: int = 5, show_routing: bool = True) -> Tuple[str, str]:
-        """
-        Generate text with comprehensive error handling and routing information
-        
-        Returns:
-            Tuple of (generated_text, routing_info_display)
-        """
+        """Generate text with comprehensive error handling and routing information"""
         start_time = time.time()
         
         # Update statistics
@@ -514,7 +454,7 @@ This analysis provides a foundation for informed decision-making while acknowled
                 # Real model generation
                 response = self._generate_real(prompt, max_length, temperature, top_p, num_encoders)
             else:
-                # Simulated generation with sophisticated responses
+                # Simulated generation
                 response = self._simulate_generation(prompt, routing_info, max_length)
             
             # Calculate performance metrics
@@ -546,46 +486,127 @@ This analysis provides a foundation for informed decision-making while acknowled
     
     def _generate_real(self, prompt: str, max_length: int, temperature: float, 
                       top_p: float, num_encoders: int) -> str:
-        """Generate using real model"""
+        """Generate using real pretrained model"""
         try:
             # Encode input
             inputs = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
             
-            # Adjust number of active encoders
+            # Adjust number of active encoders (if supported)
             if hasattr(self.model, 'set_active_encoders'):
-                self.model.set_active_encoders(min(num_encoders, self.config.max_mamba_encoders))
+                max_encoders = getattr(self.config, 'max_mamba_encoders', 100)
+                self.model.set_active_encoders(min(num_encoders, max_encoders))
             
             # Generate with memory optimization
             with torch.no_grad():
-                outputs = self.model.generate(
-                    inputs,
-                    max_length=min(max_length, getattr(self.config, 'max_sequence_length', 2048)),
-                    temperature=temperature,
-                    top_p=top_p,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id,
-                    use_cache=True
-                )
+                try:
+                    outputs = self.model.generate(
+                        inputs,
+                        max_new_tokens=min(max_length, 512),  # Limit for stability
+                        temperature=temperature,
+                        top_p=top_p,
+                        do_sample=True,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        eos_token_id=self.tokenizer.eos_token_id,
+                        use_cache=True,
+                        attention_mask=torch.ones_like(inputs)  # Ensure attention mask
+                    )
+                except Exception as gen_error:
+                    logger.warning(f"Generation with parameters failed: {gen_error}")
+                    # Fallback to simpler generation
+                    outputs = self.model.generate(
+                        inputs,
+                        max_new_tokens=min(max_length, 256),
+                        do_sample=False,  # Use greedy decoding as fallback
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        eos_token_id=self.tokenizer.eos_token_id
+                    )
             
             # Decode output
             generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             
             # Remove input prompt from output
-            response = generated_text[len(prompt):].strip()
+            if generated_text.startswith(prompt):
+                response = generated_text[len(prompt):].strip()
+            else:
+                response = generated_text.strip()
             
             return response if response else "Generated response was empty."
             
         except torch.cuda.OutOfMemoryError:
             logger.error("CUDA out of memory during generation")
-            return "Error: GPU memory insufficient. Try reducing max_length or num_encoders."
+            return "Error: GPU memory insufficient. Try reducing max_length or switching to CPU mode."
         except Exception as e:
             logger.error(f"Real generation error: {e}")
-            return f"Generation error: {str(e)}"
+            return f"Generation error: {str(e)}. Using pretrained model in fallback mode."
+    
+    def _simulate_generation(self, prompt: str, routing_info: Dict, max_length: int) -> str:
+        """Generate sophisticated simulated responses"""
+        domain = routing_info['detected_domain']
+        
+        # Enhanced domain-specific responses
+        if domain == 'code':
+            return f"""Here's a comprehensive solution for your request:
+
+```python
+def solution(input_data):
+    \"\"\"
+    Optimized implementation based on your requirements
+    \"\"\"
+    try:
+        # Input validation
+        if not input_data:
+            raise ValueError("Input cannot be empty")
+        
+        # Process the data
+        result = process_input(input_data)
+        
+        return result
+    except Exception as e:
+        print(f"Error: {{e}}")
+        return None
+
+def process_input(data):
+    # Implementation here
+    return processed_data
+```
+
+This solution includes error handling, input validation, and follows best practices for production code."""
+        
+        elif domain == 'medical':
+            return f"""Based on current medical knowledge regarding your query:
+
+**Overview:**
+This topic involves several important medical considerations that should be evaluated by healthcare professionals.
+
+**Key Points:**
+• Symptoms and presentation can vary significantly between individuals
+• Early detection and proper diagnosis are crucial
+• Treatment approaches should be personalized
+• Regular monitoring may be recommended
+
+**Important Note:** This information is for educational purposes only. Please consult with qualified healthcare professionals for personalized medical advice, diagnosis, and treatment recommendations."""
+        
+        else:
+            return f"""**Response to: "{prompt[:50]}..."**
+
+Based on analysis from {routing_info['total_active']} specialized encoders in the {domain} domain:
+
+This is a comprehensive response that addresses your query with relevant information and insights. The analysis considers multiple perspectives and provides a balanced view of the topic.
+
+**Key insights:**
+• The topic involves several interconnected factors
+• Current understanding is based on established principles
+• Practical applications may vary depending on context
+• Further exploration could yield additional insights
+
+**Domain expertise applied:** {domain.title()} specialization with {routing_info['domain_confidence']:.1%} confidence."""
     
     def _create_routing_display(self, routing_info: Dict, generation_time: float, 
                               estimated_tokens: int) -> str:
         """Create rich routing information display"""
+        model_type = "Real Pretrained Model" if (self.model_loaded and not self.fallback_mode and self.using_pretrained) else "Custom Swarm Model" if (self.model_loaded and not self.fallback_mode) else "Simulation Mode"
+        model_name = getattr(self.pretrained_loader, 'model_name', 'Custom/Simulation') if self.pretrained_loader else 'Custom/Simulation'
+        
         return f"""
 ## 🧠 Intelligent Routing Analysis
 
@@ -594,10 +615,11 @@ This analysis provides a foundation for informed decision-making while acknowled
 - **Confidence**: {routing_info['domain_confidence']:.1%}
 - **Specialization Level**: {'High' if routing_info['domain_confidence'] > 0.7 else 'Medium' if routing_info['domain_confidence'] > 0.4 else 'General'}
 
-**⚡ Encoder Activation:**
-- **Active Encoders**: {routing_info['total_active']}/{self.config.max_mamba_encoders}
-- **Selection Strategy**: Domain-optimized routing
-- **Load Distribution**: Balanced across specialized encoders
+**⚡ Model Information:**
+- **Model Type**: {model_type}
+- **Base Model**: {model_name}
+- **Active Encoders**: {routing_info['total_active']}/{getattr(self.config, 'max_mamba_encoders', 100)}
+- **Device**: {self.device}
 
 **🔢 Selected Encoder IDs:**
 {', '.join(map(str, routing_info['selected_encoders'][:15]))}{'...' if len(routing_info['selected_encoders']) > 15 else ''}
@@ -606,15 +628,15 @@ This analysis provides a foundation for informed decision-making while acknowled
 - **Generation Time**: {generation_time:.2f}s
 - **Estimated Tokens**: {estimated_tokens}
 - **Tokens/Second**: {estimated_tokens/generation_time:.1f}
-- **Model Mode**: {'Real Model' if self.model_loaded and not self.fallback_mode else 'Simulation'}
+- **Success Rate**: {(self.stats['successful_generations'] / max(self.stats['total_requests'], 1) * 100):.1f}%
 
 **🎚️ Confidence Scores (Top 5):**
 {', '.join([f'{score:.3f}' for score in routing_info['confidence_scores'][:5]])}{'...' if len(routing_info['confidence_scores']) > 5 else ''}
 
 **💡 Optimization Notes:**
 - Encoder selection optimized for domain: {routing_info['detected_domain']}
+- {'Pretrained weights from HuggingFace' if self.using_pretrained else 'Custom swarm implementation' if self.model_loaded and not self.fallback_mode else 'Simulation mode active'}
 - Dynamic load balancing across {routing_info['total_active']} active encoders
-- Confidence-weighted aggregation applied
 """
     
     def get_model_info(self) -> str:
@@ -628,21 +650,39 @@ This analysis provides a foundation for informed decision-making while acknowled
         if torch.cuda.is_available():
             gpu_info = f"{torch.cuda.get_device_name(0)} ({torch.cuda.get_device_properties(0).total_memory // 1024**3}GB)"
         
+        # Get pretrained model info if available
+        pretrained_info = ""
+        if self.pretrained_loader:
+            model_info = self.pretrained_loader.get_model_info()
+            if model_info and 'error' not in model_info:
+                pretrained_info = f"""
+**🤗 Pretrained Model Details:**
+- **Model Name**: {model_info['name']}
+- **Parameters**: {model_info['parameters']} ({model_info['parameters_millions']})
+- **Vocabulary Size**: {model_info['vocab_size']:,}
+- **Hidden Size**: {model_info['hidden_size']}
+- **Model Device**: {model_info['device']}
+- **Data Type**: {model_info['dtype']}
+"""
+        
+        status_emoji = "✅" if self.model_loaded and not self.fallback_mode else "⚠️" 
+        status_text = f"Loaded {'with Pretrained Weights' if self.using_pretrained else 'with Custom Swarm'}" if self.model_loaded and not self.fallback_mode else "Simulation Mode"
+        
         return f"""
 **🤖 Mamba Encoder Swarm Model Information**
 
 **Model Configuration:**
-- **Status**: {'✅ Loaded' if self.model_loaded else '⚠️ Simulation Mode'}
+- **Status**: {status_emoji} {status_text}
 - **Active Encoders**: {getattr(self.model, 'num_active_encoders', 'N/A')}
-- **Max Encoders**: {self.config.max_mamba_encoders}
-- **Model Dimension**: {self.config.d_model}
-- **Vocabulary Size**: {self.config.vocab_size:,}
+- **Max Encoders**: {getattr(self.config, 'max_mamba_encoders', 100)}
+- **Model Dimension**: {getattr(self.config, 'd_model', getattr(self.config, 'hidden_size', 768))}
+- **Vocabulary Size**: {getattr(self.config, 'vocab_size', 50257):,}
 - **Max Sequence Length**: {getattr(self.config, 'max_sequence_length', 'N/A')}
-
+{pretrained_info}
 **System Information:**
 - **Device**: {self.device} {f'({gpu_info})' if gpu_info != 'N/A' else ''}
 - **RAM Usage**: {memory_info.percent:.1f}% ({memory_info.used // 1024**3}GB / {memory_info.total // 1024**3}GB)
-- **Python/PyTorch**: {torch.__version__}
+- **PyTorch Version**: {torch.__version__}
 
 **Performance Statistics:**
 - **Total Requests**: {self.stats['total_requests']}
@@ -652,23 +692,83 @@ This analysis provides a foundation for informed decision-making while acknowled
 - **Avg Generation Time**: {self.stats['avg_generation_time']:.2f}s
 - **Total Tokens Generated**: {self.stats['total_tokens_generated']:,}
 
-**Fallback Mode**: {'⚠️ Active' if self.fallback_mode else '✅ Disabled'}
+**Mode**: {'🟢 Pretrained Model Active' if self.using_pretrained else '🔵 Custom Swarm Active' if self.model_loaded and not self.fallback_mode else '🟡 Simulation Mode'}
 """
     
     def get_system_status(self) -> Dict[str, Any]:
         """Get system status for monitoring"""
         return {
             'model_loaded': self.model_loaded,
+            'using_pretrained': self.using_pretrained,
             'fallback_mode': self.fallback_mode,
             'device': str(self.device),
             'stats': self.stats.copy(),
             'timestamp': datetime.now().isoformat()
         }
+    
+    def switch_model(self, model_size: str = "auto") -> str:
+        """Switch between different pretrained model sizes"""
+        if not self.using_pretrained:
+            return "❌ Model switching only available when using pretrained models"
+        
+        try:
+            MODEL_OPTIONS = {
+                "small": "state-spaces/mamba-130m",
+                "medium": "state-spaces/mamba-790m", 
+                "large": "state-spaces/mamba-1.4b",
+                "xl": "state-spaces/mamba-2.8b"
+            }
+            
+            if model_size == "auto":
+                # Auto-select based on memory
+                memory_gb = psutil.virtual_memory().total / (1024**3)
+                if memory_gb >= 32 and torch.cuda.is_available():
+                    model_size = "xl"
+                elif memory_gb >= 16 and torch.cuda.is_available():
+                    model_size = "large"
+                elif memory_gb >= 8:
+                    model_size = "medium"
+                else:
+                    model_size = "small"
+            
+            if model_size not in MODEL_OPTIONS:
+                return f"❌ Invalid model size. Choose from: {list(MODEL_OPTIONS.keys())}"
+            
+            selected_model = MODEL_OPTIONS[model_size]
+            
+            # Check if already using this model
+            if self.pretrained_loader and self.pretrained_loader.model_name == selected_model:
+                return f"✅ Already using {selected_model}"
+            
+            logger.info(f"🔄 Switching to model: {selected_model}")
+            
+            # Clear current model
+            if self.model:
+                del self.model
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+            
+            # Load new model
+            self.pretrained_loader = MambaWeightLoader(selected_model)
+            
+            if self.pretrained_loader.download_and_load():
+                self.model = self.pretrained_loader.model
+                self.tokenizer = self.pretrained_loader.tokenizer
+                self.config = self.pretrained_loader.config
+                
+                logger.info(f"✅ Successfully switched to {selected_model}")
+                return f"✅ Successfully switched to {selected_model}"
+            else:
+                logger.error(f"❌ Failed to switch to {selected_model}")
+                return f"❌ Failed to switch to {selected_model}"
+                
+        except Exception as e:
+            logger.error(f"Error switching model: {e}")
+            return f"❌ Error switching model: {str(e)}"
 
 def create_production_demo() -> gr.Blocks:
-    """Create production-ready Gradio interface"""
+    """Create production-ready Gradio interface with pretrained model support"""
     
-    # Initialize demo with fallback capability
+    # Initialize demo with pretrained model capability
     try:
         demo_instance = MambaSwarmDemo(model_path="./", fallback_mode=False)
     except Exception as e:
@@ -684,9 +784,13 @@ def create_production_demo() -> gr.Blocks:
     def refresh_model_info():
         return demo_instance.get_model_info()
     
+    def switch_model_size(model_size):
+        result = demo_instance.switch_model(model_size)
+        return result, demo_instance.get_model_info()
+    
     # Create interface
     with gr.Blocks(
-        title="Mamba Encoder Swarm - Production Demo",
+        title="Mamba Encoder Swarm - Production Demo with Pretrained Weights",
         theme=gr.themes.Soft(),
         css="""
         .gradio-container {
@@ -705,6 +809,13 @@ def create_production_demo() -> gr.Blocks:
             padding: 15px;
             margin: 10px 0;
         }
+        .status-indicator {
+            background-color: #d4edda;
+            border: 1px solid #c3e6cb;
+            border-radius: 8px;
+            padding: 10px;
+            margin: 10px 0;
+        }
         """
     ) as demo:
         
@@ -712,18 +823,29 @@ def create_production_demo() -> gr.Blocks:
         gr.Markdown("""
         # 🐍 Mamba Encoder Swarm - Production Demo
         
-        **Advanced Language Model with Dynamic Routing & Intelligent Encoder Selection**
+        **Advanced Language Model with Pretrained Weights & Dynamic Routing**
         
-        Experience the power of up to 100 specialized Mamba encoders with intelligent domain-aware routing, 
-        comprehensive error handling, and production-ready performance monitoring.
+        Now featuring **automatic pretrained weight loading** from HuggingFace's state-spaces Mamba models, 
+        with intelligent domain-aware routing across up to 100 specialized encoders.
         """)
         
         # Status indicator
         with gr.Row():
-            with gr.Column(scale=1):
+            with gr.Column(scale=3):
+                status_text = f"🟢 Real Pretrained Model" if demo_instance.using_pretrained else f"🔵 Custom Swarm Model" if demo_instance.model_loaded and not demo_instance.fallback_mode else "🟡 Simulation Mode"
                 status_indicator = gr.Markdown(
-                    f"**Status**: {'🟢 Real Model' if demo_instance.model_loaded and not demo_instance.fallback_mode else '🟡 Simulation Mode'}"
+                    f"**Status**: {status_text}",
+                    elem_classes=["status-indicator"]
                 )
+            with gr.Column(scale=1):
+                if demo_instance.using_pretrained:
+                    model_switch = gr.Dropdown(
+                        choices=["auto", "small", "medium", "large", "xl"],
+                        value="auto",
+                        label="🔄 Switch Model",
+                        info="Change pretrained model size"
+                    )
+                    switch_btn = gr.Button("Switch Model", variant="secondary", size="sm")
         
         with gr.Row():
             # Left column - Input and controls
@@ -803,7 +925,14 @@ def create_production_demo() -> gr.Blocks:
                     value=show_model_info(),
                     elem_classes=["model-info"]
                 )
-                refresh_info_btn = gr.Button("🔄 Refresh Info", size="sm")
+                with gr.Column(scale=1):
+                    refresh_info_btn = gr.Button("🔄 Refresh Info", size="sm")
+                    if demo_instance.using_pretrained:
+                        model_status = gr.Textbox(
+                            label="Model Switch Status",
+                            interactive=False,
+                            lines=2
+                        )
         
         # Examples section
         with gr.Accordion("💡 Example Prompts", open=True):
@@ -816,7 +945,8 @@ def create_production_demo() -> gr.Blocks:
                 ["Analyze the legal implications of AI-generated content", 350, 0.7, 0.9, 15, True],
                 ["Write a creative short story about a time-traveling scientist", 400, 0.9, 0.95, 12, True],
                 ["Develop a marketing strategy for a sustainable fashion startup", 300, 0.8, 0.9, 10, True],
-                ["How does quantum entanglement work and what are its applications?", 350, 0.6, 0.9, 15, True]
+                ["How does quantum entanglement work and what are its applications?", 350, 0.6, 0.9, 15, True],
+                ["Explain the economic impact of renewable energy adoption", 300, 0.7, 0.9, 12, True]
             ]
             
             gr.Examples(
@@ -828,6 +958,28 @@ def create_production_demo() -> gr.Blocks:
                 label="Click any example to load it"
             )
         
+        # Advanced features section
+        with gr.Accordion("🔬 Advanced Features", open=False):
+            gr.Markdown("""
+            ### 🚀 Pretrained Model Features
+            - **Automatic Model Selection**: Chooses optimal model size based on available memory
+            - **Dynamic Model Switching**: Switch between different Mamba model sizes
+            - **HuggingFace Integration**: Direct loading from state-spaces repository
+            - **Memory Optimization**: Efficient loading with half-precision and device mapping
+            
+            ### 🧠 Intelligent Routing System
+            - **Domain Detection**: Automatic classification of prompt domains
+            - **Specialized Encoders**: 100+ domain-specific encoder pools
+            - **Load Balancing**: Dynamic distribution across active encoders
+            - **Confidence Scoring**: Weighted aggregation based on encoder confidence
+            
+            ### 📊 Model Sizes Available
+            - **Small (130M)**: ~500MB, good for basic tasks
+            - **Medium (790M)**: ~3GB, balanced performance
+            - **Large (1.4B)**: ~5GB, high-quality responses  
+            - **XL (2.8B)**: ~10GB, best performance (requires 16GB+ RAM)
+            """)
+        
         # Event handlers
         generate_btn.click(
             fn=generate_response,
@@ -841,15 +993,36 @@ def create_production_demo() -> gr.Blocks:
             outputs=model_info_display
         )
         
+        # Model switching event handler (only if using pretrained)
+        if demo_instance.using_pretrained:
+            switch_btn.click(
+                fn=switch_model_size,
+                inputs=[model_switch],
+                outputs=[model_status, model_info_display]
+            )
+        
+        # Auto-refresh status on page load
+        demo.load(
+            fn=lambda: (demo_instance.get_model_info(), f"**Status**: {'🟢 Real Pretrained Model' if demo_instance.using_pretrained else '🔵 Custom Swarm Model' if demo_instance.model_loaded and not demo_instance.fallback_mode else '🟡 Simulation Mode'}"),
+            outputs=[model_info_display, status_indicator]
+        )
+        
         # Footer
         gr.Markdown("""
         ---
-        ### 🏗️ Architecture Overview
+        ### 🏗️ Enhanced Architecture Overview
+        
+        **🤗 Pretrained Integration**
+        - Direct loading from HuggingFace state-spaces Mamba models
+        - Automatic model size selection based on system resources
+        - Seamless fallback to custom swarm implementation
+        - Dynamic model switching without restart
         
         **🧠 Intelligent Routing System**
         - Domain detection based on prompt analysis
         - Dynamic encoder selection optimized for content type
         - Load balancing across specialized encoder pools
+        - Confidence-weighted response aggregation
         
         **🔧 Production Features**
         - Comprehensive error handling and fallback modes
@@ -861,7 +1034,7 @@ def create_production_demo() -> gr.Blocks:
         - **Medical & Healthcare** • **Legal & Regulatory** • **Code & Technical**
         - **Science & Research** • **Creative Writing** • **Business & Finance**
         
-        Built with ❤️ using Gradio, PyTorch, and the Mamba architecture
+        Built with ❤️ using Gradio, PyTorch, HuggingFace Transformers, and the Mamba architecture
         """)
     
     return demo
@@ -918,4 +1091,4 @@ if __name__ == "__main__":
             demo.launch(share=False, debug=False)
         except Exception as e2:
             logger.error(f"Minimal launch also failed: {e2}")
-            print(f"❌ All launch attempts failed. Error: {e2}")
+            print(f"❌ All launch attempts failed. Error: {e2}")