Update app.py

app.py

@@ -1,13 +1,22 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-st.set_page_config(page_title="Niche AI", layout="centered")
-
-st.title("🧠 Niche AI (CPU Test)")
-st.caption("HF Free Space · 2B params · slow but real")
+import threading
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer
+)
 
+# ---------------- CONFIG ----------------
 MODEL_ID = "Qwen/Qwen3-4B-Instruct-2507"
+MAX_NEW_TOKENS = 256
+TEMPERATURE = 0.7
+TOP_P = 0.9
+# ----------------------------------------
+
+st.set_page_config(page_title="Niche AI", layout="centered")
+st.title("🧠 Niche AI")
+st.caption("HF Free Space · CPU · Streaming")
 
 @st.cache_resource
 def load_model():
@@ -15,50 +24,71 @@ def load_model():
         MODEL_ID,
         trust_remote_code=True
     )
+
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         torch_dtype=torch.float32,
-        device_map=None   # 👈 IMPORTANT
+        device_map=None
     )
+
     model.to("cpu")
     model.eval()
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-# Session chat history
+# -------- SESSION STATE --------
 if "history" not in st.session_state:
     st.session_state.history = []
 
-prompt = st.text_input("You", placeholder="Say something...")
+# -------- INPUT --------
+prompt = st.text_input("You", placeholder="Say something…")
 
-if st.button("Send"):
-    if prompt.strip():
-        st.session_state.history.append(("You", prompt))
+if st.button("Send") and prompt.strip():
+    st.session_state.history.append(("You", prompt))
 
-        chat = [{"role": "user", "content": prompt}]
-        inputs = tokenizer.apply_chat_template(
-            chat, 
-            add_generation_prompt=True, 
-             return_tensors="pt",
-             return_dict=True   # ← this is the key
-)
+    chat = [{"role": "user", "content": prompt}]
+
+    inputs = tokenizer.apply_chat_template(
+        chat,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True
+    )
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=True,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id,
+        streamer=streamer
+    )
+
+    thread = threading.Thread(
+        target=model.generate,
+        kwargs=gen_kwargs
+    )
+    thread.start()
+
+    placeholder = st.empty()
+    output_text = ""
+
+    for token in streamer:
+        output_text += token
+        placeholder.markdown(f"**Niche:** {output_text}")
+
+    st.session_state.history.append(("Niche", output_text))
 
-        with torch.no_grad():
-            output = model.generate(
-                **inputs,
-                max_new_tokens=64,
-                do_sample=True,
-                temperature=0.8,
-                top_p=0.95,
-                eos_token_id=tokenizer.eos_token_id,   # ← add this
-                pad_token_id=tokenizer.eos_token_id
-            )
-
-        reply = tokenizer.decode(output[0], skip_special_tokens=True)
-        st.session_state.history.append(("Niche", reply))
-
-# Display chat
+# -------- DISPLAY HISTORY --------
 for speaker, text in st.session_state.history:
     if speaker == "You":
         st.markdown(f"**You:** {text}")