Update app.py

app.py

@@ -1,28 +1,58 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+st.set_page_config(page_title="Niche AI", layout="centered")
 
-st.title("Niche AI Chat")
-st.write("Chat with your trained Niche model from Hugging Face!")
+st.title("🧠 Niche AI (CPU Test)")
+st.caption("HF Free Space · 2B params · slow but real")
+
+MODEL_ID = "Neon-AI/Niche"
 
-# Load model once
 @st.cache_resource
 def load_model():
-    model_name = "Neon-AI/Niche"  # Your HF repo
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,
+        device_map=None   # 👈 IMPORTANT
+    )
+    model.to("cpu")
+    model.eval()
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-# User input
-prompt = st.text_input("You:", "")
+# Session chat history
+if "history" not in st.session_state:
+    st.session_state.history = []
+
+prompt = st.text_input("You", placeholder="Say something...")
 
 if st.button("Send"):
-    if prompt.strip() == "":
-        st.warning("Please type something!")
+    if prompt.strip():
+        st.session_state.history.append(("You", prompt))
+
+        inputs = tokenizer(prompt, return_tensors="pt")
+
+        with torch.no_grad():
+            output = model.generate(
+                **inputs,
+                max_new_tokens=64,     # keep it sane on CPU
+                do_sample=True,
+                temperature=0.8,
+                top_p=0.95
+            )
+
+        reply = tokenizer.decode(output[0], skip_special_tokens=True)
+        st.session_state.history.append(("Niche", reply))
+
+# Display chat
+for speaker, text in st.session_state.history:
+    if speaker == "You":
+        st.markdown(f"**You:** {text}")
     else:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.7)
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        st.text_area("Niche:", value=response, height=200)
+        st.markdown(f"**Niche:** {text}")