Update app.py

app.py

@@ -2,130 +2,129 @@ 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
+@st.cache_resource(show_spinner=False)
 def load_model_and_tokenizer():
     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"
-            
+        model_name = "Qwen/Qwen2.5-Coder-3B-Instruct"
+        
+        with st.spinner("🔄 Loading tokenizer..."):
             # 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
+            )
             
-            # Determine device and display info
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            st.info(f"💻 Using device: {device}")
+            # Force CPU mode and eval mode
+            model = model.to("cpu").eval()
             
-            # 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
+            # Clear memory after loading
+            gc.collect()
+            torch.cuda.empty_cache() if torch.cuda.is_available() else None
             
+            st.session_state.model_loaded = True
             return tokenizer, model
+            
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
-        raise e
+        return None, None
 
-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"""
+def generate_response(prompt, model, tokenizer, max_length=256):
     try:
-        # Tokenize input
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        # Clear memory before generation
+        gc.collect()
         
-        # Generate response with progress bar
-        with torch.no_grad(), st.spinner("🤔 Thinking..."):
+        # 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)"):
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
+                max_new_tokens=max_length,
+                temperature=0.7,
+                top_p=0.9,
                 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
+                num_beams=1,  # Disable beam search
+                early_stopping=True
             )
         
-        # Decode and return response
+        # Clear memory after generation
+        gc.collect()
+        
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response[len(prompt):].strip()
-    
+        
     except torch.cuda.OutOfMemoryError:
-        st.error("💾 GPU memory exceeded. Try reducing the maximum length or clearing the conversation.")
+        st.error("💾 Memory exceeded. Try reducing the maximum length.")
         return None
     except Exception as e:
-        st.error(f"❌ Error generating response: {str(e)}")
+        st.error(f"❌ Error: {str(e)}")
         return None
 
 # Main UI
 st.title("💬 Qwen2.5-Coder Chat")
 
-# Sidebar settings
+# Sidebar with minimal settings
 with st.sidebar:
     st.header("⚙️ Settings")
     
-    # Model settings
     max_length = st.slider(
-        "Maximum Length 📏",
+        "Response Length 📏",
         min_value=64,
-        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
+        max_value=512,
+        value=256,
+        step=64,
+        help="Shorter lengths are recommended for CPU"
     )
     
-    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
-try:
+if not st.session_state.model_loaded:
+    tokenizer, model = load_model_and_tokenizer()
+    if model is None:
+        st.stop()
+else:
     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:
@@ -148,19 +147,14 @@ if prompt := st.chat_input("💭 Ask me anything about coding..."):
     
     # Generate and display response
     with st.chat_message("assistant"):
-        # 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:"
+        # Keep only last message for context to reduce memory usage
+        conversation = f"Human: {prompt}\nAssistant:"
         
         response = generate_response(
             conversation,
             model,
             tokenizer,
-            max_new_tokens=max_length,
-            temperature=temperature,
-            top_p=top_p
+            max_length=max_length
         )
         
         if response:
@@ -174,4 +168,7 @@ if prompt := st.chat_input("💭 Ask me anything about coding..."):
                 "timestamp": timestamp
             })
         else:
-            st.error("❌ Failed to generate response. Please try again with different settings.")
+            st.error("❌ Failed to generate response. Please try again with a shorter length.")
+            
+        # Clear memory after response
+        gc.collect()