Update app.py

app.py

@@ -9,9 +9,9 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 @st.cache_resource
 def load_model():
-    model_id = "tiiuae/falcon-7b-instruct"
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+    model_id = "sshleifer/tiny-gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 generator = load_model()
@@ -25,9 +25,9 @@ TRIGGER_PHRASES = ["kill myself", "end it all", "suicide", "not worth living", "
 def is_high_risk(text):
     return any(phrase in text.lower() for phrase in TRIGGER_PHRASES)
 
-def get_response(prompt, max_new_tokens=150, temperature=0.7):
+def get_response(prompt, max_new_tokens=100, temperature=0.7):
     output = generator(prompt, max_new_tokens=max_new_tokens, temperature=temperature)[0]["generated_text"]
-    return output.split("AI:")[-1].strip() if "AI:" in output else output.strip()
+    return output.strip()
 
 st.title("🧠 TARS.help")
 st.markdown("### A quiet AI that reflects and replies.")
@@ -56,7 +56,7 @@ if st.button("🧾 Generate Session Summary"):
     convo = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history])
     prompt = f"Summarize the emotional tone and themes in this conversation:\n{convo}\nSummary:"
     try:
-        summary = get_response(prompt, max_new_tokens=200, temperature=0.5)
+        summary = get_response(prompt, max_new_tokens=100, temperature=0.5)
         st.session_state.summary = summary
     except Exception as e:
         st.error("Summary generation failed.")