Update app.py

app.py

@@ -22,7 +22,7 @@ st.title("🚀 WizNerd Insp 🚀")
 # Sidebar configuration
 with st.sidebar:
     st.header("Authentication 🔒")
-    hf_token = st.text_input("HuggingFace Token", type="password", 
+    hf_token = st.text_input("Hugging Face Token", type="password", 
                            help="Get your token from https://huggingface.co/settings/tokens")
     
     if not hf_token:
@@ -90,7 +90,7 @@ def generate_response(prompt, file_context):
     Answer:"""
     
     # Create streamer
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
     # Tokenize input
     inputs = tokenizer(
@@ -107,7 +107,9 @@ def generate_response(prompt, file_context):
         max_new_tokens=1024,
         temperature=0.7,
         top_p=0.9,
-        repetition_penalty=1.1
+        repetition_penalty=1.1,
+        do_sample=True,
+        use_cache=True  # Enable caching for faster generation
     )
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
@@ -160,12 +162,10 @@ if prompt := st.chat_input("Ask your inspection question..."):
             with st.chat_message("assistant", avatar="🤖"):
                 streamer = generate_response(prompt, file_context)
                 response = st.write_stream(streamer)
+                # Remove <think> tokens if present
+                response = response.replace("<think>", "").replace("</think>", "").strip()
                 st.session_state.messages.append({"role": "assistant", "content": response})
-        except:
-            # Fallback if avatar fails
-            with st.chat_message("assistant"):
-                streamer = generate_response(prompt, file_context)
-                response = st.write_stream(streamer)
-                st.session_state.messages.append({"role": "assistant", "content": response})
+        except Exception as e:
+            st.error(f"⚡ Generation error: {str(e)}")
     else:
         st.error("🤖 Model not loaded - check your token and connection!")