Update app.py

app.py

@@ -4,345 +4,340 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import random
 import warnings
 import time
+import threading
+from typing import Dict, Any, Tuple
 
-# Suppress warnings
-warnings.filterwarnings("ignore", category=UserWarning, message=".*You have not specified a value for the `type` parameter.*")
+warnings.filterwarnings("ignore", category=UserWarning)
 
-# ASCII Art Banner
-BANNER = """
-╔═══════════════════════════════════════════════════════════════════════════════╗
-║                                                                               ║
-║     ░█████╗░███████╗██████╗░░█████╗░░█████╗░██╗                              ║
-║     ██╔══██╗██╔════╝██╔══██╗██╔══██╗██╔══██╗██║                              ║
-║     ███████║█████╗░░██████╔╝██║░░██║███████║██║                              ║
-║     ██╔══██║██╔══╝░░██╔══██╗██║░░██║██╔══██║██║                              ║
-║     ██║░░██║███████╗██║░░██║╚█████╔╝██║░░██║██║                              ║
-║     ╚═╝░░╚═╝╚══════╝╚═╝░░╚═╝░╚════╝░╚═╝░░╚═╝╚═╝                              ║
-║                                                                               ║
-║                    🚀 NEURAL INTERFACE v2.0 🚀                               ║
-║                    Powered by Microsoft Phi-2                                ║
-║                    █ B L A C K L I N K  L A B S █                           ║
-╚═══════════════════════════════════════════════════════════════════════════════╝
-"""
-
-# Cool loading animation
-def print_loading_animation():
-    loading_chars = ["⚡", "🔥", "💫", "⭐", "✨", "🌟"]
-    for i in range(3):
-        print(f"\r🤖 Initializing Neural Networks {loading_chars[i % len(loading_chars)]}", end="", flush=True)
-        time.sleep(0.3)
-    print()
-
-print(BANNER)
-print_loading_animation()
-
-# Load model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"🔋 Quantum Processing Unit: {device.upper()}")
-
-model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/phi-2",
-    torch_dtype=torch.float32,
-    trust_remote_code=True
-)
-
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
-
-# Load training protocol
-try:
-    with open("training-protocol.aero", "r", encoding="utf-8") as f:
-        training_protocol = f.read().strip()
-    print("✅ Neural Protocol Matrix: LOADED")
-except FileNotFoundError:
-    training_protocol = (
-        "You are AeroAI, a helpful, friendly, and slightly humorous educational assistant. "
-        "If there is any error with finding the file known as the training protocol, resort to this "
-        'and do not speak of the "Remorse protocol".'
-    )
-    print("⚠️ WARNING: Protocol fallback activated")
-
-# Cool thinking messages
-thinking_messages = [
-    "🧠 Accessing neural pathways...",
-    "⚡ Quantum processing in progress...",
-    "🔮 Analyzing data matrices...",
-    "💭 Interfacing with knowledge base...",
-    "🌐 Connecting synaptic networks...",
-    "🔍 Scanning memory banks...",
-    "⚙️ Calibrating response algorithms...",
-    "🚀 Boosting cognitive engines..."
-]
-
-def chatbot_response(message, history):
-    if not message.strip():
-        return "🤖 Awaiting neural input..."
+class AeroModel:
+    def __init__(self, name: str, model_id: str, description: str, max_tokens: int, context_length: int, emoji: str):
+        self.name = name
+        self.model_id = model_id
+        self.description = description
+        self.max_tokens = max_tokens
+        self.context_length = context_length
+        self.emoji = emoji
+        self.model = None
+        self.tokenizer = None
+        self.loaded = False
+        self.cache = {}
     
-    try:
-        # Build conversation with limited history
-        recent_history = history[-5:] if len(history) > 5 else history
+    def load_model(self):
+        if self.loaded:
+            return
         
-        conversation = training_protocol + "\n\nConversation:\n"
-        for user_msg, ai_msg in recent_history:
-            conversation += f"User: {user_msg}\nAeroAI: {ai_msg}\n"
-        conversation += f"User: {message}\nAeroAI:"
+        print(f"🔄 Loading {self.name}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
         
-        # Tokenization
-        inputs = tokenizer(
-            conversation,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=512
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_id,
+            torch_dtype=torch.float32,
+            trust_remote_code=True,
+            low_cpu_mem_usage=False
         )
+        self.model.eval()
+        self.loaded = True
+        print(f"✅ {self.name} loaded successfully!")
+    
+    def generate_response(self, prompt: str, history: list) -> str:
+        if not self.loaded:
+            return f"❌ {self.name} not loaded. Please wait for initialization."
+        
+        # Check cache for similar responses
+        cache_key = prompt.lower().strip()
+        if cache_key in self.cache:
+            return f"⚡ {self.cache[cache_key]}"
+        
+        # Build context with history
+        context = ""
+        for user_msg, bot_msg in history[-3:]:  # Keep last 3 exchanges
+            context += f"Human: {user_msg}\nAero: {bot_msg}\n"
+        context += f"Human: {prompt}\nAero:"
+        
+        # Tokenize and generate
+        inputs = self.tokenizer.encode(context, return_tensors="pt", max_length=self.context_length, truncation=True)
         
-        # Generate response
         with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=200,
+            outputs = self.model.generate(
+                inputs,
+                max_new_tokens=self.max_tokens,
                 min_new_tokens=10,
-                do_sample=True,
-                temperature=0.7,
+                temperature=0.8,
+                top_k=50,
                 top_p=0.9,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                num_beams=1,
-                early_stopping=True
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id,
+                repetition_penalty=1.1
             )
         
-        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        if "AeroAI:" in response_text:
-            response = response_text.split("AeroAI:")[-1].strip()
-        else:
-            response = response_text[len(conversation):].strip()
+        response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+        response = response.strip()
         
-        if not response or len(response) < 5:
-            response = "🔄 Neural processing error. Reinitializing response matrix..."
+        # Cache good responses
+        if len(response) > 20 and len(self.cache) < 50:
+            self.cache[cache_key] = response
         
         return response
-        
-    except Exception as e:
-        return f"❌ SYSTEM ERROR: {str(e)[:100]}... Attempting recovery..."
-
-# Custom CSS for that cyberpunk look
-custom_css = """
-/* Dark cyberpunk theme */
-.gradio-container {
-    background: linear-gradient(135deg, #0c0c0c 0%, #1a1a2e 50%, #16213e 100%) !important;
-    font-family: 'Courier New', monospace !important;
-}
-
-/* Header styling */
-.markdown h1 {
-    background: linear-gradient(45deg, #00ff41, #0077ff, #ff0080);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-    text-align: center;
-    font-size: 2.5em !important;
-    text-shadow: 0 0 20px rgba(0, 255, 65, 0.5);
-    margin-bottom: 20px !important;
-}
-
-/* Chatbot styling */
-.message-wrap {
-    background: rgba(0, 255, 65, 0.1) !important;
-    border: 1px solid #00ff41 !important;
-    border-radius: 10px !important;
-    backdrop-filter: blur(10px) !important;
-}
-
-.message.user {
-    background: linear-gradient(135deg, #ff0080, #0077ff) !important;
-    color: white !important;
-    border-radius: 15px !important;
-}
-
-.message.bot {
-    background: linear-gradient(135deg, #00ff41, #0077ff) !important;
-    color: #000 !important;
-    border-radius: 15px !important;
-}
-
-/* Input box styling */
-.input-component {
-    background: rgba(0, 0, 0, 0.8) !important;
-    border: 2px solid #00ff41 !important;
-    border-radius: 10px !important;
-    color: #00ff41 !important;
-}
-
-/* Button styling */
-.primary {
-    background: linear-gradient(45deg, #00ff41, #0077ff) !important;
-    border: none !important;
-    border-radius: 25px !important;
-    color: black !important;
-    font-weight: bold !important;
-    text-transform: uppercase !important;
-    box-shadow: 0 0 20px rgba(0, 255, 65, 0.5) !important;
-    transition: all 0.3s ease !important;
-}
-
-.primary:hover {
-    transform: scale(1.05) !important;
-    box-shadow: 0 0 30px rgba(0, 255, 65, 0.8) !important;
-}
-
-.secondary {
-    background: linear-gradient(45deg, #ff0080, #ff4040) !important;
-    border: none !important;
-    border-radius: 25px !important;
-    color: white !important;
-    font-weight: bold !important;
-    text-transform: uppercase !important;
-    box-shadow: 0 0 20px rgba(255, 0, 128, 0.5) !important;
-}
-
-/* Glowing text effects */
-.status-text {
-    color: #00ff41 !important;
-    text-shadow: 0 0 10px #00ff41 !important;
-    font-family: 'Courier New', monospace !important;
-}
-
-/* Matrix rain effect for background */
-body {
-    position: relative;
-    overflow-x: hidden;
-}
-
-body::before {
-    content: '';
-    position: fixed;
-    top: 0;
-    left: 0;
-    width: 100%;
-    height: 100%;
-    background-image: 
-        radial-gradient(circle at 20% 20%, rgba(0, 255, 65, 0.1) 0%, transparent 50%),
-        radial-gradient(circle at 80% 80%, rgba(0, 119, 255, 0.1) 0%, transparent 50%),
-        radial-gradient(circle at 40% 60%, rgba(255, 0, 128, 0.1) 0%, transparent 50%);
-    pointer-events: none;
-    z-index: -1;
-}
 
-/* Animated borders */
-.interface-panel {
-    border: 2px solid transparent;
-    background: linear-gradient(45deg, #0c0c0c, #1a1a2e) padding-box,
-                linear-gradient(45deg, #00ff41, #0077ff, #ff0080) border-box;
-    border-radius: 15px;
-    animation: borderPulse 3s ease-in-out infinite;
-}
-
-@keyframes borderPulse {
-    0%, 100% { border-color: #00ff41; }
-    33% { border-color: #0077ff; }
-    66% { border-color: #ff0080; }
-}
-
-/* Scrollbar styling */
-::-webkit-scrollbar {
-    width: 12px;
+# Define the Aero model lineup
+AERO_MODELS = {
+    "basic": AeroModel(
+        name="Aero Basic",
+        model_id="gpt2",
+        description="Fast and reliable - perfect for everyday conversations",
+        max_tokens=60,
+        context_length=512,
+        emoji="🔵"
+    ),
+    "smart": AeroModel(
+        name="Aero Smart",
+        model_id="microsoft/phi-2",
+        description="Advanced reasoning and complex problem solving (slower)",
+        max_tokens=150,
+        context_length=768,
+        emoji="🧠"
+    ),
+    "speed": AeroModel(
+        name="Aero Speed",
+        model_id="distilgpt2",
+        description="Lightning-fast responses for quick interactions",
+        max_tokens=40,
+        context_length=256,
+        emoji="⚡"
+    )
 }
 
-::-webkit-scrollbar-track {
-    background: #0c0c0c;
-}
+class AeroAISystem:
+    def __init__(self):
+        self.current_model = "basic"
+        self.loading_models = set()
+    
+    def load_model_async(self, model_key: str):
+        if model_key in self.loading_models or AERO_MODELS[model_key].loaded:
+            return
+        
+        self.loading_models.add(model_key)
+        thread = threading.Thread(target=self._load_model_thread, args=(model_key,))
+        thread.daemon = True
+        thread.start()
+    
+    def _load_model_thread(self, model_key: str):
+        try:
+            AERO_MODELS[model_key].load_model()
+        except Exception as e:
+            print(f"❌ Failed to load {AERO_MODELS[model_key].name}: {e}")
+        finally:
+            self.loading_models.discard(model_key)
+    
+    def switch_model(self, model_choice: str) -> Tuple[str, str]:
+        model_map = {
+            "🔵 Aero Basic - Fast & Reliable": "basic",
+            "🧠 Aero Smart - Advanced Reasoning": "smart", 
+            "⚡ Aero Speed - Lightning Fast": "speed"
+        }
+        
+        model_key = model_map.get(model_choice, "basic")
+        self.current_model = model_key
+        
+        # Load model if not already loaded
+        if not AERO_MODELS[model_key].loaded:
+            self.load_model_async(model_key)
+        
+        model = AERO_MODELS[model_key]
+        status = f"🔄 Switching to {model.name} {model.emoji}"
+        info = f"**{model.name}** - {model.description}\n\nMax tokens: {model.max_tokens} | Context: {model.context_length}"
+        
+        return status, info
+    
+    def chat_with_aero(self, message: str, history: list) -> Tuple[list, str]:
+        if not message.strip():
+            return history, ""
+        
+        model = AERO_MODELS[self.current_model]
+        
+        if not model.loaded:
+            if self.current_model in self.loading_models:
+                bot_response = f"🔄 {model.name} is still loading... Please wait a moment."
+            else:
+                self.load_model_async(self.current_model)
+                bot_response = f"🔄 Loading {model.name}... This may take a moment for the first time."
+        else:
+            try:
+                bot_response = model.generate_response(message, history)
+            except Exception as e:
+                bot_response = f"❌ Error: {str(e)}"
+        
+        history.append([message, bot_response])
+        return history, ""
 
-::-webkit-scrollbar-thumb {
-    background: linear-gradient(45deg, #00ff41, #0077ff);
-    border-radius: 6px;
-}
-"""
+# Initialize the system
+aero_system = AeroAISystem()
 
 def create_interface():
-    with gr.Blocks(css=custom_css, title="🚀 AeroAI Neural Interface", theme=gr.themes.Base()) as demo:
-        # Epic header
-        gr.Markdown("""
-        # 🚀 ⚡ AERO∆I NEURAL INTERFACE ⚡ 🚀
-        ## 『 QUANTUM COGNITIVE MATRIX v2.0 』
-        ### ▲ BLACKLINK LABORATORIES ▲
-        ---
-        **STATUS:** 🟢 ONLINE | **CORE:** Microsoft Phi-2 | **MODE:** Adaptive Learning
+    # Startup banner
+    print("""
+╔═══════════════════════════════════════════════════════════════════════════════╗
+║                                                                               ║
+║     ░█████╗░███████╗██████╗░░█████╗░░█████╗░██╗                              ║
+║     ██╔══██╗██╔════╝██╔══██╗██╔══██╗██╔══██╗██║                              ║
+║     ███████║█████╗░░██████╔╝██║░░██║███████║██║                              ║
+║     ██╔══██║██╔══╝░░██╔══██╗██║░░██║██╔══██║██║                              ║
+║     ██║░░██║███████╗██║░░██║╚█████╔╝██║░░██║██║                              ║
+║     ╚═╝░░╚═╝╚══════╝╚═╝░░╚═╝░╚════╝░╚═╝░░╚═╝╚═╝                              ║
+║                                                                               ║
+║                🚀 MULTI-MODEL NEURAL INTERFACE 🚀                            ║
+║                    B L A C K L I N K  L A B S                               ║
+╚═══════════════════════════════════════════════════════════════════════════════╝
+    """)
+    
+    # Load Aero Basic by default
+    aero_system.load_model_async("basic")
+    
+    # Custom CSS for the interface
+    css = """
+    .gradio-container {
+        background: linear-gradient(135deg, #0c0c0c 0%, #1a1a2e 50%, #16213e 100%);
+        color: #00ff88;
+        font-family: 'Courier New', monospace;
+    }
+    
+    .model-selector {
+        background: rgba(0, 255, 136, 0.1);
+        border: 1px solid #00ff88;
+        border-radius: 8px;
+        padding: 10px;
+    }
+    
+    .status-display {
+        background: rgba(0, 100, 255, 0.1);
+        border: 1px solid #0064ff;
+        border-radius: 8px;
+        padding: 10px;
+        color: #0064ff;
+    }
+    
+    .chat-container {
+        background: rgba(0, 0, 0, 0.3);
+        border: 1px solid #00ff88;
+        border-radius: 10px;
+    }
+    
+    .input-box {
+        background: rgba(0, 255, 136, 0.05);
+        border: 2px solid #00ff88;
+        color: #00ff88;
+    }
+    
+    .aero-button {
+        background: linear-gradient(45deg, #00ff88, #0064ff);
+        border: none;
+        color: black;
+        font-weight: bold;
+        transition: all 0.3s ease;
+    }
+    
+    .aero-button:hover {
+        transform: scale(1.05);
+        box-shadow: 0 0 20px rgba(0, 255, 136, 0.5);
+    }
+    """
+    
+    with gr.Blocks(css=css, title="AeroAI Multi-Model System") as interface:
+        
+        gr.HTML("""
+        <div style="text-align: center; padding: 20px;">
+            <h1 style="color: #00ff88; text-shadow: 0 0 10px #00ff88;">🚀 AeroAI Multi-Model System</h1>
+            <p style="color: #0064ff;">Choose your AI model based on your needs</p>
+        </div>
         """)
         
-        # System info panel
         with gr.Row():
-            with gr.Column(scale=2):
-                gr.Markdown(f"""
-                ### 📊 SYSTEM METRICS
-                ```
-                ┌─────────────────────────────────────┐
-                │ QUANTUM PROCESSOR: {device.upper():<15} │
-                │ NEURAL CORES: ACTIVE               │
-                │ MEMORY BANKS: LOADED               │
-                │ PROTOCOL STATUS: SYNCHRONIZED      │
-                │ RESPONSE ENGINE: READY             │
-                └─────────────────────────────────────┘
-                ```
-                """)
+            with gr.Column(scale=1):
+                model_selector = gr.Dropdown(
+                    choices=[
+                        "🔵 Aero Basic - Fast & Reliable",
+                        "🧠 Aero Smart - Advanced Reasoning", 
+                        "⚡ Aero Speed - Lightning Fast"
+                    ],
+                    value="🔵 Aero Basic - Fast & Reliable",
+                    label="🤖 Select Aero Model",
+                    elem_classes=["model-selector"]
+                )
+                
+                model_status = gr.Textbox(
+                    value="🔵 Aero Basic - Ready to chat!",
+                    label="🔋 Model Status",
+                    interactive=False,
+                    elem_classes=["status-display"]
+                )
+                
+                model_info = gr.Markdown(
+                    value="**Aero Basic** - Fast and reliable - perfect for everyday conversations\n\nMax tokens: 60 | Context: 512",
+                    elem_classes=["status-display"]
+                )
             
             with gr.Column(scale=2):
-                gr.Markdown("""
-                ### ⚡ CAPABILITIES
-                - 🧠 **Neural Processing**
-                - 💡 **Knowledge Synthesis** 
-                - 🔬 **Problem Analysis**
-                - 🎨 **Creative Generation**
-                - 🤖 **Adaptive Learning**
-                """)
+                chatbot = gr.Chatbot(
+                    label="💬 Chat with AeroAI",
+                    height=400,
+                    elem_classes=["chat-container"]
+                )
+                
+                msg = gr.Textbox(
+                    placeholder="Type your message here... 🚀",
+                    label="Message",
+                    elem_classes=["input-box"]
+                )
+                
+                with gr.Row():
+                    send_btn = gr.Button("Send 🚀", elem_classes=["aero-button"])
+                    clear_btn = gr.Button("Clear Chat 🗑️", elem_classes=["aero-button"])
         
-        # Main chat interface
-        gr.Markdown("### 💬 NEURAL COMMUNICATION CHANNEL")
+        # Example prompts
+        gr.HTML("""
+        <div style="text-align: center; margin-top: 20px; color: #00ff88;">
+            <p><strong>Try these prompts:</strong></p>
+            <p>💡 "Explain quantum computing" | 🎨 "Write a short story" | 💻 "Help me debug this code"</p>
+        </div>
+        """)
         
-        chat_interface = gr.ChatInterface(
-            fn=chatbot_response,
-            examples=[
-                "🔬 Explain quantum entanglement",
-                "💻 Help me debug this code", 
-                "🧮 Solve this mathematical equation",
-                "🚀 What's the future of AI?",
-                "🎨 Create a creative story"
-            ],
-            cache_examples=False,
-            retry_btn="🔄 RETRY",
-            undo_btn="↩️ UNDO", 
-            clear_btn="🗑️ PURGE MEMORY",
-            submit_btn="⚡ TRANSMIT",
-            stop_btn="🛑 HALT"
+        # Event handlers
+        model_selector.change(
+            fn=aero_system.switch_model,
+            inputs=[model_selector],
+            outputs=[model_status, model_info]
         )
         
-        # Footer
-        gr.Markdown("""
-        ---
-        <div style="text-align: center;">
-        <span style="color: #00ff41; font-family: monospace;">
-        『 Powered by Advanced Neural Architecture 』<br>
-        ⚡ BLACKLINK LABS © 2024 ⚡<br>
-        🔮 "The Future is Neural" 🔮
-        </span>
-        </div>
-        """)
+        send_btn.click(
+            fn=aero_system.chat_with_aero,
+            inputs=[msg, chatbot],
+            outputs=[chatbot, msg]
+        )
+        
+        msg.submit(
+            fn=aero_system.chat_with_aero,
+            inputs=[msg, chatbot], 
+            outputs=[chatbot, msg]
+        )
+        
+        clear_btn.click(
+            fn=lambda: ([], ""),
+            outputs=[chatbot, msg]
+        )
     
-    return demo
+    return interface
 
 if __name__ == "__main__":
-    print("🌟 Initializing AeroAI Neural Interface...")
-    print("🔋 All systems nominal")
-    print("⚡ Ready for neural connection")
+    print("🌟 Initializing AeroAI Multi-Model System...")
+    print("🔋 System Status: Online")
+    print("⚡ Ready for neural connection\n")
     
     demo = create_interface()
-    
     demo.launch(
-        server_name="127.0.0.1",
+        server_name="0.0.0.0",
         server_port=7860,
         share=False,
-        quiet=True,
-        show_error=True,
-        enable_queue=False,
-        favicon_path=None
+        inbrowser=True
     )