Upload app.py

app.py

@@ -40,65 +40,149 @@ class UltimateModelLoader:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
         # Comprehensive model configurations
-        self.model_configs = {
-            # Reliable models (priority 1-3)
+        self.model_configs = self._get_all_available_models()
+        
+    def _get_all_available_models(self):
+        """Get all available models including trained checkpoints"""
+        models = {}
+        
+        # Check for custom trained models first (highest priority)
+        trained_models = self._discover_trained_models()
+        for model_name, config in trained_models.items():
+            models[model_name] = config
+        
+        # Standard models with adjusted priorities
+        models.update({
+            # Priority Mamba models - adjusted priorities for trained models
+            "state-spaces/mamba-130m": {
+                "display_name": "Mamba 130M Encoder",
+                "size": "small",
+                "priority": 10,  # Lower priority than trained models
+                "reliable": True,
+                "params": 130_000_000,
+                "vocab_size": 50280,
+                "d_model": 768
+            },
+            "state-spaces/mamba-790m": {
+                "display_name": "Mamba 790M Encoder",
+                "size": "large",
+                "priority": 11,
+                "reliable": True,
+                "params": 790_000_000,
+                "vocab_size": 50280,
+                "d_model": 1536
+            },
+            "state-spaces/mamba-1.4b": {
+                "display_name": "Mamba 1.4B Encoder",
+                "size": "xlarge",
+                "priority": 12,
+                "reliable": True,
+                "params": 1_400_000_000,
+                "vocab_size": 50280,
+                "d_model": 2048
+            },
+            # Fallback models (priority 20-27) - Only used if Mamba fails
             "gpt2-medium": {
-                "display_name": "GPT2 Medium (355M)",
+                "display_name": "GPT2 Medium (355M) [Fallback]",
                 "size": "medium",
-                "priority": 1,
+                "priority": 20,
                 "reliable": True,
                 "params": 355_000_000
             },
             "gpt2": {
-                "display_name": "GPT2 Base (117M)", 
+                "display_name": "GPT2 Base (117M) [Fallback]", 
                 "size": "small",
-                "priority": 2,
+                "priority": 21,
                 "reliable": True,
                 "params": 117_000_000
             },
             "distilgpt2": {
-                "display_name": "DistilGPT2 (82M)",
+                "display_name": "DistilGPT2 (82M) [Fallback]",
                 "size": "small",
-                "priority": 3,
+                "priority": 22,
                 "reliable": True,
                 "params": 82_000_000
             },
-            # Advanced models (priority 4-7)
             "microsoft/DialoGPT-medium": {
-                "display_name": "DialoGPT Medium (355M)",
+                "display_name": "DialoGPT Medium (355M) [Fallback]",
                 "size": "medium",
-                "priority": 4,
+                "priority": 23,
                 "reliable": True,
                 "params": 355_000_000
-            },
-            "state-spaces/mamba-130m": {
-                "display_name": "Mamba 130M",
-                "size": "small",
-                "priority": 5,
-                "reliable": False,  # Needs validation
-                "params": 130_000_000,
-                "vocab_size": 50280,
-                "d_model": 768
-            },
-            "state-spaces/mamba-790m": {
-                "display_name": "Mamba 790M",
-                "size": "large",
-                "priority": 6,
-                "reliable": False,
-                "params": 790_000_000,
-                "vocab_size": 50280,
-                "d_model": 1536
-            },
-            "state-spaces/mamba-1.4b": {
-                "display_name": "Mamba 1.4B",
-                "size": "xlarge",
-                "priority": 7,
-                "reliable": False,
-                "params": 1_400_000_000,
-                "vocab_size": 50280,
-                "d_model": 2048
             }
-        }
+        })
+        
+        return models
+    
+    def _discover_trained_models(self):
+        """Discover custom trained models in checkpoints directory"""
+        trained_models = {}
+        
+        # Check for checkpoint directories
+        checkpoint_dirs = [
+            "checkpoints",
+            "mamba_checkpoints", 
+            "training_output"
+        ]
+        
+        priority = 1  # Highest priority for trained models
+        
+        for checkpoint_dir in checkpoint_dirs:
+            if os.path.exists(checkpoint_dir):
+                for item in os.listdir(checkpoint_dir):
+                    item_path = os.path.join(checkpoint_dir, item)
+                    
+                    # Check if it's a model directory with config.json
+                    config_path = os.path.join(item_path, "config.json")
+                    if os.path.isdir(item_path) and os.path.exists(config_path):
+                        
+                        try:
+                            import json
+                            with open(config_path, 'r') as f:
+                                model_config = json.load(f)
+                            
+                            # Estimate model size from config
+                            d_model = model_config.get('d_model', model_config.get('hidden_size', 768))
+                            n_layers = model_config.get('n_layers', model_config.get('num_hidden_layers', 12))
+                            vocab_size = model_config.get('vocab_size', 50257)
+                            
+                            # Estimate parameters
+                            estimated_params = d_model * d_model * n_layers * 4  # Rough estimate
+                            
+                            # Determine size category
+                            if estimated_params < 200_000_000:
+                                size = "small"
+                            elif estimated_params < 800_000_000:
+                                size = "medium"
+                            elif estimated_params < 1_500_000_000:
+                                size = "large"
+                            else:
+                                size = "xlarge"
+                            
+                            trained_models[item_path] = {
+                                "display_name": f"🎯 Custom Trained: {item} ({d_model}D)",
+                                "size": size,
+                                "priority": priority,
+                                "reliable": True,
+                                "params": estimated_params,
+                                "vocab_size": vocab_size,
+                                "d_model": d_model,
+                                "is_custom": True,
+                                "local_path": item_path
+                            }
+                            
+                            priority += 1
+                            
+                        except Exception as e:
+                            logger.warning(f"Could not load config for {item_path}: {e}")
+                            continue
+        
+        if trained_models:
+            logger.info(f"🎯 Found {len(trained_models)} custom trained models!")
+            for name, config in trained_models.items():
+                logger.info(f"  - {config['display_name']}")
+        
+        return trained_models
         
         # Generation configurations by model size
         self.generation_configs = {
@@ -1209,18 +1293,22 @@ def create_ultimate_interface():
     ) as demo:
         
         gr.Markdown("""
-        # 🐍 Mamba Encoder Swarm v1.0
+        # 🐍 Ultimate Mamba Encoder Swarm - Production Intelligence System
+        
+        **🚀 Advanced AI Language Model with True Mamba Encoder Swarm Intelligence**
         
-        **🚀 Advanced AI Language Model with Mamba Swarm Intelligence**
+        Features cutting-edge **Mamba State-Space Models**, advanced domain routing, comprehensive performance analytics, and multi-tier quality protection.
         
-        Features cutting-edge model selection, advanced domain routing, comprehensive performance analytics, and multi-tier quality protection.
+        **🔥 Now Prioritizing REAL Mamba Encoders over GPT2 fallbacks!**
         """)
         
         # Ultimate status display
         with gr.Row():
-            status_text = "🟢 Ultimate AI System Online" if swarm.model_loaded else "🟡 Protected Fallback Mode"
-            model_info = f" | Model: {swarm.model_loader.model_name} ({swarm.current_model_size.title()})" if swarm.model_loaded else ""
-            gr.Markdown(f"**System Status**: {status_text}{model_info}", elem_classes=["status-box"])
+            status_text = "🟢 Mamba Encoder System Online" if swarm.model_loaded else "🟡 Protected Fallback Mode"
+            model_info = f" | Active: {swarm.model_loader.model_name} ({swarm.current_model_size.title()})" if swarm.model_loaded else ""
+            is_mamba = "mamba" in swarm.model_loader.model_name.lower() if swarm.model_loaded and swarm.model_loader.model_name else False
+            encoder_type = "🐍 MAMBA ENCODERS" if is_mamba else "⚠️ FALLBACK MODE"
+            gr.Markdown(f"**{encoder_type}**: {status_text}{model_info}", elem_classes=["status-box"])
         
         with gr.Row():
             # Ultimate control panel
@@ -1310,15 +1398,17 @@ def create_ultimate_interface():
         # Ultimate footer
         gr.Markdown("""
         ---
-        ### 🌟 Ultimate Production Features
-        - **🧠 Advanced Model Intelligence** - Dynamic model selection with size control (Small/Medium/Large/XLarge)
+        ### 🐍 True Mamba Encoder Swarm Features
+        - **🧠 Real Mamba State-Space Models** - Prioritized Mamba-130M, Mamba-790M, Mamba-1.4B encoders
         - **🎯 Elite Domain Routing** - 7 specialized domains with confidence-based encoder selection  
-        - **⚡ GPU Acceleration** - Optimized CUDA operations with memory management
+        - **⚡ Advanced State-Space Processing** - Leveraging Mamba's selective state-space architecture
         - **🛡️ Zero-Gibberish Guarantee** - Multi-layer quality validation prevents nonsense output
         - **📊 Ultimate Analytics** - Real-time performance monitoring with comprehensive metrics
-        - **🔄 Smart Fallbacks** - Advanced multi-tier fallback protection system
-        - **🎛️ Dynamic Control** - Real-time model switching and parameter optimization
-        - **🚀 Production Ready** - Enterprise-grade reliability and error handling
+        - **🔄 Smart Fallbacks** - GPT2 models only used if Mamba encoders fail to load
+        - **🎛️ Dynamic Control** - Real-time model switching between different Mamba sizes
+        - **🚀 Production Ready** - Enterprise-grade reliability with true encoder swarm intelligence
+        
+        **Note**: System prioritizes Mamba encoders over traditional transformers for authentic swarm behavior!
         """)
     
     return demo