Update app.py

app.py

@@ -1,192 +1,197 @@
 import streamlit as st
 import os
-import time
-from model_manager import ModelManager
-from inference_engine import InferenceEngine
 import torch
+import warnings
+import logging
+from transformers import AutoTokenizer
+import gc
 
-# Page configuration
+# HF Spaces specific configuration
 st.set_page_config(
-    page_title="Automotive SLM Chatbot",
+    page_title="🚗 Automotive SLM Assistant",
     page_icon="🚗",
-    layout="wide",
-    initial_sidebar_state="expanded"
+    layout="wide"
 )
 
-# Custom CSS
-st.markdown("""
-<style>
-.main-header {
-    font-size: 2.5rem;
-    color: #1f77b4;
-    text-align: center;
-    margin-bottom: 2rem;
-}
-.chat-message {
-    padding: 1rem;
-    border-radius: 0.5rem;
-    margin: 0.5rem 0;
-}
-.user-message {
-    background-color: #e3f2fd;
-    border-left: 4px solid #1976d2;
-}
-.assistant-message {
-    background-color: #f3e5f5;
-    border-left: 4px solid #7b1fa2;
-}
-.model-info {
-    background-color: #f5f5f5;
-    padding: 1rem;
-    border-radius: 0.5rem;
-    border: 1px solid #ddd;
-}
-</style>
-""", unsafe_allow_html=True)
+# Suppress warnings for HF Spaces
+warnings.filterwarnings('ignore')
+logging.getLogger('streamlit').setLevel(logging.ERROR)
+logging.getLogger('transformers').setLevel(logging.ERROR)
 
-@st.cache_resource
-def load_model_manager():
-    """Cache the model manager to avoid reloading"""
-    return ModelManager("assets")
+# 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 initialize_session_state():
-    """Initialize session state variables"""
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-    if "current_model" not in st.session_state:
-        st.session_state.current_model = None
-    if "inference_engine" not in st.session_state:
-        st.session_state.inference_engine = None
+# 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 display_chat_message(role, content, model_info=None):
-    """Display a chat message with proper styling"""
-    if role == "user":
-        st.markdown(f"""
-        <div class="chat-message user-message">
-            <strong>You:</strong> {content}
-        </div>
-        """, unsafe_allow_html=True)
-    else:
-        model_text = f" ({model_info})" if model_info else ""
-        st.markdown(f"""
-        <div class="chat-message assistant-message">
-            <strong>Assistant{model_text}:</strong> {content}
-        </div>
-        """, unsafe_allow_html=True)
+def generate_response(model, tokenizer, prompt, max_tokens=50, temperature=0.8):
+    """Generate response optimized for HF Spaces"""
+    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."
+        
+    except Exception as e:
+        return f"I encountered an error: {str(e)}. Please try again."
 
 def main():
-    # Initialize session state
-    initialize_session_state()
+    # Title and description
+    st.title("🚗 Automotive SLM Assistant")
+    st.markdown("*Specialized AI assistant for automotive questions and troubleshooting*")
     
-    # Header
-    st.markdown('<h1 class="main-header">🚗 Automotive SLM Chatbot</h1>', unsafe_allow_html=True)
+    # Load model
+    model, tokenizer, config = load_model_for_hf_spaces()
     
-    # Load model manager
-    model_manager = load_model_manager()
+    if model is None:
+        st.stop()
     
-    # Sidebar for model selection and settings
+    # Sidebar settings
     with st.sidebar:
-        st.header("⚙️ Model Settings")
-        
-        # Model selection
-        available_models = model_manager.get_available_models()
-        if available_models:
-            selected_model = st.selectbox(
-                "Select Model:",
-                available_models,
-                index=0 if st.session_state.current_model is None else available_models.index(st.session_state.current_model) if st.session_state.current_model in available_models else 0
-            )
-            
-            # Load model if changed
-            if selected_model != st.session_state.current_model:
-                with st.spinner(f"Loading {selected_model}..."):
-                    model, tokenizer, config = model_manager.load_model(selected_model)
-                    st.session_state.inference_engine = InferenceEngine(model, tokenizer, config)
-                    st.session_state.current_model = selected_model
-                st.success(f"Model {selected_model} loaded successfully!")
-        else:
-            st.error("No models found in assets folder!")
-            st.stop()
-        
-        # Model information
-        if st.session_state.inference_engine:
-            st.subheader("📊 Model Info")
-            model_info = model_manager.get_model_info(selected_model)
-            st.markdown(f"""
-            <div class="model-info">
-                <strong>Model:</strong> {model_info['name']}<br>
-                <strong>Type:</strong> {model_info['type']}<br>
-                <strong>Parameters:</strong> {model_info['parameters']}<br>
-                <strong>Size:</strong> {model_info['size']}
-            </div>
-            """, unsafe_allow_html=True)
+        st.header("⚙️ Settings")
+        max_tokens = st.slider("Response Length", 20, 100, 50)
+        temperature = st.slider("Creativity", 0.3, 1.5, 0.8, 0.1)
         
-        # Generation settings
-        st.subheader("🎛️ Generation Settings")
-        max_tokens = st.slider("Max Tokens", 10, 200, 50)
-        temperature = st.slider("Temperature", 0.1, 2.0, 0.8, 0.1)
-        top_p = st.slider("Top P", 0.1, 1.0, 0.9, 0.05)
-        top_k = st.slider("Top K", 1, 100, 50)
-        
-        # Clear chat button
-        if st.button("🗑️ Clear Chat"):
-            st.session_state.messages = []
-            st.rerun()
+        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
+        """)
     
-    # Main chat interface
-    if st.session_state.inference_engine is None:
-        st.info("Please select a model from the sidebar to start chatting.")
-        return
+    # 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
-    chat_container = st.container()
-    with chat_container:
-        for message in st.session_state.messages:
-            display_chat_message(
-                message["role"], 
-                message["content"], 
-                message.get("model", None)
-            )
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
     
     # Chat input
-    prompt = st.chat_input("Ask me about automotive topics...")
-    
-    if prompt:
+    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)
         
-        # Display user message
-        with chat_container:
-            display_chat_message("user", 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})
         
-        # Generate response
-        with st.spinner("Generating response..."):
-            try:
-                response = st.session_state.inference_engine.generate_response(
-                    prompt,
-                    max_tokens=max_tokens,
-                    temperature=temperature,
-                    top_p=top_p,
-                    top_k=top_k
-                )
-                
-                # Add assistant message
-                st.session_state.messages.append({
-                    "role": "assistant", 
-                    "content": response,
-                    "model": selected_model
-                })
-                
-                # Display assistant message
-                with chat_container:
-                    display_chat_message("assistant", response, selected_model)
-                    
-            except Exception as e:
-                st.error(f"Error generating response: {str(e)}")
-    
-    # Footer
-    st.markdown("---")
-    st.markdown("*Powered by Automotive SLM - Specialized for automotive assistance*")
+        # 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()
 
 if __name__ == "__main__":
     main()