Update app.py

app.py

@@ -1,129 +1,139 @@
-import os
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import datetime
 
-# Page configuration
+# Set page configuration
 st.set_page_config(
-    page_title="💬 Qwen2.5-Coder Chat",
+    page_title="Qwen2.5-Coder Chat",
     page_icon="💬",
     layout="wide"
 )
 
-# Set cache directory explicitly for Hugging Face Spaces
-os.environ["TRANSFORMERS_CACHE"] = "/root/.cache/huggingface"
-
-# Initialize session state for conversation history
+# Initialize session state
 if 'messages' not in st.session_state:
     st.session_state.messages = []
 
-# Cache model loading to prevent re-loading each session
 @st.cache_resource
 def load_model_and_tokenizer():
-    model_name = "Qwen/Qwen2.5-Coder-3B-Instruct"  # Smaller 3B model for efficiency
-
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name, 
-        trust_remote_code=True
-    )
-
-    # Device configuration
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    st.info(f"Using device: {device}")
+    try:
+        # Display loading message
+        with st.spinner("🔄 Loading model and tokenizer... This might take a few minutes..."):
+            model_name = "Qwen/Qwen2.5-Coder-3B-Instruct"
+            
+            # Load tokenizer first
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                trust_remote_code=True
+            )
+            
+            # Determine device and display info
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            st.info(f"💻 Using device: {device}")
+            
+            # Load model with appropriate settings
+            if device == "cuda":
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    torch_dtype=torch.float16,  # Use float16 for GPU
+                    device_map="auto",
+                    trust_remote_code=True
+                ).eval()  # Set to evaluation mode
+            else:
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    device_map={"": device},
+                    trust_remote_code=True,
+                    low_cpu_mem_usage=True
+                ).eval()  # Set to evaluation mode
+            
+            return tokenizer, model
+    except Exception as e:
+        st.error(f"❌ Error loading model: {str(e)}")
+        raise e
 
-    # Load model with optimizations for CPU
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float32 if device == "cpu" else torch.float16,
-        device_map="auto" if device == "cuda" else {"": device},
-        trust_remote_code=True,
-        low_cpu_mem_usage=True  # Reduce memory usage for CPU
-    )
+def generate_response(prompt, model, tokenizer, max_new_tokens=512, temperature=0.7, top_p=0.9):
+    """Generate response from the model with better error handling"""
+    try:
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        # Generate response with progress bar
+        with torch.no_grad(), st.spinner("🤔 Thinking..."):
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1,
+                no_repeat_ngram_size=3
+            )
+        
+        # Decode and return response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response[len(prompt):].strip()
     
-    return tokenizer, model
+    except torch.cuda.OutOfMemoryError:
+        st.error("💾 GPU memory exceeded. Try reducing the maximum length or clearing the conversation.")
+        return None
+    except Exception as e:
+        st.error(f"❌ Error generating response: {str(e)}")
+        return None
 
-# Title
+# Main UI
 st.title("💬 Qwen2.5-Coder Chat")
 
 # Sidebar settings
 with st.sidebar:
-    st.header("Settings")
+    st.header("⚙️ Settings")
     
+    # Model settings
     max_length = st.slider(
-        "Maximum Length",
+        "Maximum Length 📏",
         min_value=64,
-        max_value=1024,  # Lowered for CPU
-        value=256,  # Default setting for CPU
-        step=64,
-        help="Maximum number of tokens to generate"
+        max_value=2048,
+        value=512,
+        step=64
     )
     
     temperature = st.slider(
-        "Temperature",
+        "Temperature 🌡️",
         min_value=0.1,
-        max_value=1.5,  # Lower range to make output more deterministic
-        value=0.5,
-        step=0.1,
-        help="Higher values make output more random, lower values more deterministic"
+        max_value=2.0,
+        value=0.7,
+        step=0.1
     )
     
     top_p = st.slider(
-        "Top P",
+        "Top P 📊",
         min_value=0.1,
         max_value=1.0,
-        value=0.8,
-        step=0.1,
-        help="Nucleus sampling: higher values consider more tokens, lower values are more focused"
+        value=0.9,
+        step=0.1
     )
     
-    if st.button("Clear Conversation"):
+    # Clear conversation button
+    if st.button("🗑️ Clear Conversation"):
         st.session_state.messages = []
         st.rerun()
 
-# Load model with caching
+# Load model
 try:
-    with st.spinner("Loading model... Please wait..."):
-        tokenizer, model = load_model_and_tokenizer()
+    tokenizer, model = load_model_and_tokenizer()
 except Exception as e:
-    st.error(f"Error loading model: {str(e)}")
+    st.error("❌ Failed to load model. Please check the logs and refresh the page.")
     st.stop()
 
-# Response generation function
-def generate_response(prompt, max_new_tokens=256, temperature=0.5, top_p=0.8):
-    """Generate response from the model"""
-    try:
-        # Tokenize the input
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        
-        # Generate response
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=True,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-            )
-        
-        # Decode and return response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response[len(prompt):].strip()  # Extract only the model's response
-    
-    except Exception as e:
-        st.error(f"Error generating response: {str(e)}")
-        return None
-
 # Display conversation history
-for message in st.session_state.messages[-5:]:  # Limit to last 5 messages for efficiency
+for message in st.session_state.messages:
     with st.chat_message(message["role"]):
-        st.write(f"{message['content']}\n\n_{message['timestamp']}_")
+        st.markdown(f"{message['content']}\n\n_{message['timestamp']}_")
 
 # Chat input
-if prompt := st.chat_input("Ask me anything about coding..."):
+if prompt := st.chat_input("💭 Ask me anything about coding..."):
     # Add user message
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     st.session_state.messages.append({
@@ -134,31 +144,34 @@ if prompt := st.chat_input("Ask me anything about coding..."):
     
     # Display user message
     with st.chat_message("user"):
-        st.write(f"{prompt}\n\n_{timestamp}_")
+        st.markdown(f"{prompt}\n\n_{timestamp}_")
     
     # Generate and display response
     with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            # Prepare conversation context, limited to recent exchanges
-            conversation = "\n".join(
-                f"{'Human' if msg['role'] == 'user' else 'Assistant'}: {msg['content']}" 
-                for msg in st.session_state.messages[-3:]  # Send only the last 3 messages
-            ) + "\nAssistant:"
-            
-            response = generate_response(
-                conversation,
-                max_new_tokens=max_length,
-                temperature=temperature,
-                top_p=top_p
-            )
+        # Prepare conversation context (limit to last 3 messages to prevent context overflow)
+        conversation = "\n".join(
+            f"{'Human' if msg['role'] == 'user' else 'Assistant'}: {msg['content']}" 
+            for msg in st.session_state.messages[-3:]
+        ) + "\nAssistant:"
+        
+        response = generate_response(
+            conversation,
+            model,
+            tokenizer,
+            max_new_tokens=max_length,
+            temperature=temperature,
+            top_p=top_p
+        )
+        
+        if response:
+            timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            st.markdown(f"{response}\n\n_{timestamp}_")
             
-            if response:
-                timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                st.write(f"{response}\n\n_{timestamp}_")
-                
-                # Add response to chat history
-                st.session_state.messages.append({
-                    "role": "assistant",
-                    "content": response,
-                    "timestamp": timestamp
-                })
+            # Add response to chat history
+            st.session_state.messages.append({
+                "role": "assistant",
+                "content": response,
+                "timestamp": timestamp
+            })
+        else:
+            st.error("❌ Failed to generate response. Please try again with different settings.")