Update app.py

app.py

@@ -1,197 +1,78 @@
-import streamlit as st
 import os
-import torch
-import warnings
-import logging
-from transformers import AutoTokenizer
-import gc
+import gradio as gr
+from src.model_manager import ModelManager
+from src.inference_engine import InferenceEngine
 
-# HF Spaces specific configuration
-st.set_page_config(
-    page_title="🚗 Automotive SLM Assistant",
-    page_icon="🚗",
-    layout="wide"
-)
+ASSETS_DIR = "assets"
 
-# Suppress warnings for HF Spaces
-warnings.filterwarnings('ignore')
-logging.getLogger('streamlit').setLevel(logging.ERROR)
-logging.getLogger('transformers').setLevel(logging.ERROR)
+# Initialize once
+manager = ModelManager(os.path.join(ASSETS_DIR, "models"))
 
-# HF Spaces optimized model loading
-@st.cache_resource(show_spinner="🚀 Loading your Automotive AI Assistant...")
-def load_model_for_hf_spaces():
-    """Optimized model loading for HF Spaces environment"""
-    try:
-        # Force CPU usage for HF Spaces
-        device = torch.device('cpu')
-        
-        # Load tokenizer first
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        
-        # Simplified model loading for HF Spaces
-        # You would upload your model files to the HF Spaces repo
-        if os.path.exists("model.pt"):
-            checkpoint = torch.load("model.pt", map_location=device)
-            
-            # Create simple config if not in checkpoint
-            config = {
-                'd_model': 256,
-                'n_layer': 4,
-                'n_head': 4,
-                'vocab_size': 50257,
-                'n_positions': 256,
-                'use_moe': True,
-                'n_experts': 4
-            }
-            
-            # Use simplified model class for HF Spaces
-            model = SimpleAutomotiveModel(config)
-            
-            if 'model_state_dict' in checkpoint:
-                model.load_state_dict(checkpoint['model_state_dict'])
-            
-            model.eval()
-            
-            return model, tokenizer, config
-        else:
-            st.error("Model file not found. Please upload your model.pt to the repository.")
-            return None, None, None
-            
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        return None, None, None
+def list_models():
+    models = manager.get_available_models()
+    return models
+
+# Cache loaded engines by model name
+_engines = {}
 
-# Simplified model class for HF Spaces
-class SimpleAutomotiveModel(torch.nn.Module):
-    """Simplified model for HF Spaces deployment"""
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.embeddings = torch.nn.Embedding(config['vocab_size'], config['d_model'])
-        self.layers = torch.nn.ModuleList([
-            torch.nn.TransformerEncoderLayer(
-                d_model=config['d_model'], 
-                nhead=config['n_head'],
-                batch_first=True
-            ) for _ in range(config['n_layer'])
-        ])
-        self.ln_f = torch.nn.LayerNorm(config['d_model'])
-        self.lm_head = torch.nn.Linear(config['d_model'], config['vocab_size'], bias=False)
-    
-    def forward(self, input_ids):
-        x = self.embeddings(input_ids)
-        for layer in self.layers:
-            x = layer(x)
-        x = self.ln_f(x)
-        return {"logits": self.lm_head(x)}
-    
-    def generate(self, input_ids, max_new_tokens=50, temperature=0.8, **kwargs):
-        """Simple generation for HF Spaces"""
-        device = input_ids.device
-        generated = input_ids.clone()
-        
-        for _ in range(max_new_tokens):
-            with torch.no_grad():
-                outputs = self.forward(generated)
-                logits = outputs["logits"][:, -1, :] / temperature
-                probs = torch.softmax(logits, dim=-1)
-                next_token = torch.multinomial(probs, 1)
-                generated = torch.cat([generated, next_token], dim=1)
-                
-                # Simple stopping condition
-                if next_token.item() == 50256:  # EOS token
-                    break
-        
-        return generated
+def load_engine(model_name):
+    if model_name in _engines:
+        return _engines[model_name]
+    model, tokenizer, config = manager.load_model(model_name)
+    engine = InferenceEngine(model, tokenizer, config)
+    _engines[model_name] = engine
+    return engine
 
-def generate_response(model, tokenizer, prompt, max_tokens=50, temperature=0.8):
-    """Generate response optimized for HF Spaces"""
+def chat_fn(message, history, model_name, max_tokens, temperature, top_p, top_k):
     try:
-        # Tokenize
-        inputs = tokenizer(prompt, return_tensors="pt", max_length=200, truncation=True)
-        
-        # Generate
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs['input_ids'],
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                pad_token_id=tokenizer.pad_token_id
-            )
-        
-        # Decode
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Remove original prompt
-        if response.startswith(prompt):
-            response = response[len(prompt):].strip()
-        
-        return response if response else "I apologize, but I couldn't generate a proper response. Please try rephrasing your question."
-        
+        engine = load_engine(model_name)
     except Exception as e:
-        return f"I encountered an error: {str(e)}. Please try again."
+        return history + [[message, f"Error loading model: {e}"]]
+    reply = engine.generate_response(
+        message,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k
+    )
+    history = history + [[message, reply]]
+    return history
+
+def clear_chat():
+    return []
 
-def main():
-    # Title and description
-    st.title("🚗 Automotive SLM Assistant")
-    st.markdown("*Specialized AI assistant for automotive questions and troubleshooting*")
-    
-    # Load model
-    model, tokenizer, config = load_model_for_hf_spaces()
-    
-    if model is None:
-        st.stop()
-    
-    # Sidebar settings
-    with st.sidebar:
-        st.header("⚙️ Settings")
-        max_tokens = st.slider("Response Length", 20, 100, 50)
-        temperature = st.slider("Creativity", 0.3, 1.5, 0.8, 0.1)
-        
-        st.markdown("---")
-        st.markdown("### 🎯 Example Questions")
-        st.markdown("""
-        - How do I check tire pressure?
-        - What does the check engine light mean?
-        - How to jump start a car?
-        - Electric vehicle charging tips
-        - Brake maintenance schedule
-        """)
-    
-    # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = [
-            {"role": "assistant", "content": "Hello! I'm your Automotive AI Assistant. Ask me anything about cars, maintenance, troubleshooting, or automotive technology!"}
-        ]
-    
-    # Display chat history
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-    
-    # Chat input
-    if prompt := st.chat_input("Ask me about automotive topics..."):
-        # Add user message
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.markdown(prompt)
-        
-        # Generate and display response
-        with st.chat_message("assistant"):
-            with st.spinner("🤔 Thinking..."):
-                response = generate_response(model, tokenizer, prompt, max_tokens, temperature)
-                st.markdown(response)
-                st.session_state.messages.append({"role": "assistant", "content": response})
-        
-        # Cleanup for HF Spaces memory management
-        if len(st.session_state.messages) > 20:  # Keep last 20 messages
-            st.session_state.messages = st.session_state.messages[-20:]
-        
-        # Force garbage collection
-        gc.collect()
+with gr.Blocks(title="Automotive SLM Chatbot") as demo:
+    gr.Markdown("# 🚗 Automotive SLM Chatbot (Gradio)")
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(height=450, label="Chat")
+            msg = gr.Textbox(placeholder="Ask about automotive topics...", label="Your message")
+            with gr.Row():
+                send_btn = gr.Button("Send", variant="primary")
+                clear_btn = gr.Button("Clear")
+        with gr.Column(scale=2):
+            gr.Markdown("### Model settings")
+            available = list_models()
+            model_dropdown = gr.Dropdown(
+                choices=available, value=available[0] if available else None, label="Model"
+            )
+            max_tokens = gr.Slider(10, 256, value=64, step=1, label="Max tokens")
+            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+            top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")
+            gr.Markdown("Tip: lower temperature for more deterministic answers.")
+    # Events
+    send_evt = send_btn.click(
+        fn=chat_fn,
+        inputs=[msg, chatbot, model_dropdown, max_tokens, temperature, top_p, top_k],
+        outputs=[chatbot]
+    )
+    msg.submit(
+        fn=chat_fn,
+        inputs=[msg, chatbot, model_dropdown, max_tokens, temperature, top_p, top_k],
+        outputs=[chatbot]
+    )
+    clear_btn.click(clear_chat, inputs=None, outputs=[chatbot])
 
 if __name__ == "__main__":
-    main()
+    demo.launch()