Update app.py

app.py

@@ -3,73 +3,98 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 from threading import Thread
 import torch
 
-# 1. Page Configuration (No Sidebar)
-st.set_page_config(page_title="Claude Clone", page_icon="🤖", layout="centered")
+# 1. Page Configuration
+st.set_page_config(page_title="Qwen Chat", page_icon="🧠", layout="centered")
 
+# Custom CSS for a cleaner "Claude-like" feel
 st.markdown("""
     <style>
         [data-testid="stSidebar"] {display: none;}
-        .stChatMessage {border-radius: 15px; padding: 10px; margin-bottom: 10px;}
+        .stChatMessage { border-radius: 10px; margin-bottom: 5px; }
+        .stChatInputContainer { padding-bottom: 20px; }
     </style>
 """, unsafe_allow_html=True)
 
-st.title("Qwen 2.5 Coder 1.5B 🚀")
-st.caption("Now with real-time streaming and optimized CPU inference.")
+st.title("Qwen 2.5 3B Chat 🚀")
+st.caption("A balanced, high-performance model for local CPU/GPU inference.")
 
 # 2. Optimized Model Loading
-@st.cache_resource
+@st.cache_resourced
 def load_model():
-    model_id = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+    # '3B' is the most feasible mid-point for modern laptops/PCs
+    model_id = "Qwen/Qwen2.5-3B-Instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     
-    # Use bfloat16 for speed on modern CPUs, or float32 for maximum compatibility
+    # Auto-detect device (Use GPU if available, else CPU)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Use float16 for GPU or bfloat16 for modern CPUs to save memory
+    dtype = torch.float16 if device == "cuda" else torch.bfloat16
+    
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
-        torch_dtype=torch.float32, # CPU-friendly
+        torch_dtype=dtype,
         device_map="auto"
     )
-    return model, tokenizer
+    return model, tokenizer, device
 
-model, tokenizer = load_model()
+model, tokenizer, device = load_model()
 
-# 3. Session State
+# 3. Session State for Chat History
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display History
+# Display Chat History
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
 # 4. Chat Input & Streaming Logic
-if prompt := st.chat_input("Ask me anything..."):
+if prompt := st.chat_input("How can I help you today?"):
+    # Add user message to history
     st.session_state.messages.append({"role": "user", "content": prompt})
+    
     with st.chat_message("user"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        # Set up the streamer
+        # Setup Streamer
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         
-        # Prepare the input
-        messages = [{"role": "system", "content": "You are a helpful coding assistant."}] + st.session_state.messages
-        inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
+        # Format conversation using the model's chat template
+        # Limit history to last 5 turns to prevent CPU slowdown
+        context_messages = st.session_state.messages[-10:] 
+        full_prompt = [{"role": "system", "content": "You are Qwen, a helpful and concise AI assistant."}] + context_messages
+        
+        model_inputs = tokenizer.apply_chat_template(
+            full_prompt, 
+            tokenize=True, 
+            add_generation_prompt=True, 
+            return_tensors="pt"
+        ).to(device)
 
-        # Run generation in a separate thread to allow UI to remain responsive
+        # Generation Arguments
         generation_kwargs = dict(
-            input_ids=inputs,
+            input_ids=model_inputs,
             streamer=streamer,
-            max_new_tokens=512,
+            max_new_tokens=1024,
             do_sample=True,
             temperature=0.7,
-            top_p=0.9,
+            top_p=0.8,
+            repetition_penalty=1.1,
             pad_token_id=tokenizer.eos_token_id
         )
-        
+
+        # Start thread
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        # Stream the response to the UI
-        full_response = st.write_stream(streamer)
+        # Stream the output
+        response_container = st.empty()
+        full_response = ""
         
+        # Use st.write_stream for a native feel
+        full_response = st.write_stream(streamer)
+            
+    # Save assistant response to history
     st.session_state.messages.append({"role": "assistant", "content": full_response})