Update app.py

app.py

@@ -2,129 +2,130 @@ import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import datetime
-import gc
-import os
-
-# Enable memory efficient options
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
 
 # Set page configuration
 st.set_page_config(
     page_title="Qwen2.5-Coder Chat",
     page_icon="💬",
-    layout="wide",
+    layout="wide"
 )
 
 # Initialize session state
 if 'messages' not in st.session_state:
     st.session_state.messages = []
-if 'model_loaded' not in st.session_state:
-    st.session_state.model_loaded = False
 
-@st.cache_resource(show_spinner=False)
+@st.cache_resource
 def load_model_and_tokenizer():
     try:
-        model_name = "Qwen/Qwen2.5-Coder-3B-Instruct"
-        
-        with st.spinner("🔄 Loading tokenizer..."):
+        # 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
             )
-        
-        with st.spinner("🔄 Loading model... (this may take a few minutes on CPU)"):
-            # Load model with 8-bit quantization for CPU
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                device_map={"": "cpu"},
-                trust_remote_code=True,
-                low_cpu_mem_usage=True,
-                torch_dtype=torch.float32,
-                load_in_8bit=True  # Enable 8-bit quantization
-            )
             
-            # Force CPU mode and eval mode
-            model = model.to("cpu").eval()
+            # Determine device and display info
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            st.info(f"💻 Using device: {device}")
             
-            # Clear memory after loading
-            gc.collect()
-            torch.cuda.empty_cache() if torch.cuda.is_available() else None
+            # 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
             
-            st.session_state.model_loaded = True
             return tokenizer, model
-            
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
-        return None, None
+        raise e
 
-def generate_response(prompt, model, tokenizer, max_length=256):
+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:
-        # Clear memory before generation
-        gc.collect()
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         
-        # Tokenize with shorter maximum length
-        inputs = tokenizer(
-            prompt, 
-            return_tensors="pt",
-            max_length=512,
-            truncation=True
-        ).to("cpu")
-        
-        # Generate with minimal parameters for CPU
-        with torch.no_grad(), st.spinner("🤔 Thinking... (please be patient)"):
+        # Generate response with progress bar
+        with torch.no_grad(), st.spinner("🤔 Thinking..."):
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=max_length,
-                temperature=0.7,
-                top_p=0.9,
+                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,
-                num_beams=1,  # Disable beam search
-                early_stopping=True
+                repetition_penalty=1.1,
+                no_repeat_ngram_size=3
             )
         
-        # Clear memory after generation
-        gc.collect()
-        
+        # Decode and return response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response[len(prompt):].strip()
-        
+    
     except torch.cuda.OutOfMemoryError:
-        st.error("💾 Memory exceeded. Try reducing the maximum length.")
+        st.error("💾 GPU memory exceeded. Try reducing the maximum length or clearing the conversation.")
         return None
     except Exception as e:
-        st.error(f"❌ Error: {str(e)}")
+        st.error(f"❌ Error generating response: {str(e)}")
         return None
 
 # Main UI
 st.title("💬 Qwen2.5-Coder Chat")
 
-# Sidebar with minimal settings
+# Sidebar settings
 with st.sidebar:
     st.header("⚙️ Settings")
     
+    # Model settings
     max_length = st.slider(
-        "Response Length 📏",
+        "Maximum Length 📏",
         min_value=64,
-        max_value=512,
-        value=256,
-        step=64,
-        help="Shorter lengths are recommended for CPU"
+        max_value=2048,
+        value=512,
+        step=64
+    )
+    
+    temperature = st.slider(
+        "Temperature 🌡️",
+        min_value=0.1,
+        max_value=2.0,
+        value=0.7,
+        step=0.1
     )
     
+    top_p = st.slider(
+        "Top P 📊",
+        min_value=0.1,
+        max_value=1.0,
+        value=0.9,
+        step=0.1
+    )
+    
+    # Clear conversation button
     if st.button("🗑️ Clear Conversation"):
         st.session_state.messages = []
         st.rerun()
 
 # Load model
-if not st.session_state.model_loaded:
-    tokenizer, model = load_model_and_tokenizer()
-    if model is None:
-        st.stop()
-else:
+try:
     tokenizer, model = load_model_and_tokenizer()
+except Exception as e:
+    st.error("❌ Failed to load model. Please check the logs and refresh the page.")
+    st.stop()
 
 # Display conversation history
 for message in st.session_state.messages:
@@ -147,14 +148,19 @@ if prompt := st.chat_input("💭 Ask me anything about coding..."):
     
     # Generate and display response
     with st.chat_message("assistant"):
-        # Keep only last message for context to reduce memory usage
-        conversation = f"Human: {prompt}\nAssistant:"
+        # 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_length=max_length
+            max_new_tokens=max_length,
+            temperature=temperature,
+            top_p=top_p
         )
         
         if response:
@@ -168,7 +174,4 @@ if prompt := st.chat_input("💭 Ask me anything about coding..."):
                 "timestamp": timestamp
             })
         else:
-            st.error("❌ Failed to generate response. Please try again with a shorter length.")
-            
-        # Clear memory after response
-        gc.collect()
+            st.error("❌ Failed to generate response. Please try again with different settings.")