Update app.py

app.py

@@ -3,14 +3,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 from threading import Thread
 import torch
 
-# Clean, centered layout without sidebar
-st.set_page_config(page_title="Qwen 3 4B Stream", page_icon="⚡", layout="centered", initial_sidebar_state="collapsed")
+# UI Setup - No Sidebar
+st.set_page_config(page_title="Gemma 3 1B Fast Chat", page_icon="💎", layout="centered", initial_sidebar_state="collapsed")
 
-# 1. Model Configuration (Qwen 3 4B - 4-bit for speed)
-MODEL_ID = "unsloth/Qwen3-4B-Instruct-2507-bnb-4bit"
+# 1. Model ID for Gemma 3 1B Instruct
+MODEL_ID = "google/gemma-3-1b-it"
 
 @st.cache_resource
-def load_resource():
+def load_model():
+    # Gemma 3 1B is small enough to load in bfloat16 or float32 quickly
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
@@ -19,61 +20,57 @@ def load_resource():
     )
     return tokenizer, model
 
-tokenizer, model = load_resource()
+tokenizer, model = load_model()
 
-# Custom CSS to hide the sidebar toggle
+# Custom CSS to keep the clean look
 st.markdown("<style>[data-testid='collapsedControl'] { display: none; }</style>", unsafe_allow_html=True)
 
-st.title("⚡ Qwen 3 4B Stream")
-st.caption("Real-time local generation | No Sidebar")
+st.title("💎 Gemma 3 1B")
+st.caption("Lightweight Google AI | High-Speed Local Chat")
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Action Buttons
-col1, col2 = st.columns([5, 1])
-with col2:
-    if st.button("Reset"):
-        st.session_state.messages = []
-        st.rerun()
+# Action Button
+if st.button("Clear Chat History"):
+    st.session_state.messages = []
+    st.rerun()
 
-# Display chat history
+# Display history
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
 
 # 2. Streaming Chat Input
-if prompt := st.chat_input("Ask Qwen 3..."):
+if prompt := st.chat_input("Message Gemma 3..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        # Setup the Streamer
+        # Setup Streamer
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         
-        # Prepare input
+        # Gemma 3 uses a specific chat template format
         input_text = tokenizer.apply_chat_template(st.session_state.messages, tokenize=False, add_generation_prompt=True)
         inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
 
-        # 3. Generation in a separate thread
+        # Threaded generation for real-time streaming
         generation_kwargs = dict(
             **inputs,
             streamer=streamer,
             max_new_tokens=1024,
             do_sample=True,
             temperature=0.7,
-            top_p=0.8,
-            pad_token_id=tokenizer.eos_token_id
+            top_p=0.9,
         )
         
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        # 4. Stream to UI
+        # Update UI word-by-word
         placeholder = st.empty()
         full_response = ""
-        
         for new_text in streamer:
             full_response += new_text
             placeholder.markdown(full_response + "▌")