Update app.py

app.py

@@ -3,14 +3,23 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 from threading import Thread
 import torch
 
-# 1. Page Config - No Sidebar
-st.set_page_config(page_title="Qwen 3 0.6B Instant", page_icon="⚡", layout="centered", initial_sidebar_state="collapsed")
+# 1. Page Configuration (Centered and No Sidebar)
+st.set_page_config(
+    page_title="Qwen 3 0.6B Chat", 
+    page_icon="⚡", 
+    layout="centered", 
+    initial_sidebar_state="collapsed"
+)
 
-# Model ID: The 2026 ultra-lightweight version
-MODEL_ID = "Qwen/Qwen3-0.6B-Instruct"
+# Custom CSS to hide the sidebar toggle button entirely
+st.markdown("<style>[data-testid='collapsedControl'] { display: none; }</style>", unsafe_allow_html=True)
+
+# 2. Model & Tokenizer Initialization (Using your direct load logic)
+MODEL_ID = "Qwen/Qwen3-0.6B"
 
 @st.cache_resource
-def load_resource():
+def load_llm():
+    # Loading the tokenizer and model directly as requested
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
@@ -19,45 +28,41 @@ def load_resource():
     )
     return tokenizer, model
 
-tokenizer, model = load_resource()
-
-# Hide Sidebar Toggle
-st.markdown("<style>[data-testid='collapsedControl'] { display: none; }</style>", unsafe_allow_html=True)
+tokenizer, model = load_llm()
 
-st.title("⚡ Qwen 3 0.6B: Instant")
-st.caption("The fastest chat model of 2026 | Pure CPU Speed")
+# 3. Chat UI Logic
+st.title("⚡ Qwen 3 0.6B")
+st.caption("Using your direct-load logic with real-time streaming.")
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Header button for clearing
-if st.button("Reset Chat"):
-    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. Chat Input & Generation
-if prompt := st.chat_input("Ask me anything..."):
+# 4. Input & Streaming Generation
+if prompt := st.chat_input("Ask Qwen 3..."):
+    # Store and display user message
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
+        # Initialize the streamer
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         
-        # Apply Qwen 3 template (disabling 'thinking' for maximum chat speed)
-        input_text = tokenizer.apply_chat_template(
-            st.session_state.messages, 
-            tokenize=False, 
+        # Using your chat template logic
+        inputs = tokenizer.apply_chat_template(
+            st.session_state.messages,
             add_generation_prompt=True,
-            enable_thinking=False 
-        )
-        inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(model.device)
 
+        # Background thread for generation
         generation_kwargs = dict(
             **inputs,
             streamer=streamer,
@@ -65,14 +70,16 @@ if prompt := st.chat_input("Ask me anything..."):
             do_sample=True,
             temperature=0.7,
             top_p=0.8,
+            pad_token_id=tokenizer.eos_token_id
         )
         
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        # Word-by-word streaming
+        # Update the UI as tokens arrive
         placeholder = st.empty()
         full_response = ""
+        
         for new_text in streamer:
             full_response += new_text
             placeholder.markdown(full_response + "▌")