Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| from datetime import datetime | |
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| # -- SETUP -- | |
| os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" | |
| def load_respondent(): | |
| model_id = "microsoft/phi-2" # switched to a coherent, safe small model | |
| tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| model = AutoModelForCausalLM.from_pretrained(model_id) | |
| return pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| generator = load_respondent() | |
| if "history" not in st.session_state: | |
| st.session_state.history = [] | |
| st.session_state.summary = "" | |
| # -- STYLING -- | |
| st.markdown(""" | |
| <style> | |
| body { | |
| background-color: #111827; | |
| color: #f3f4f6; | |
| } | |
| .stTextInput > div > div > input { | |
| color: #f3f4f6; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # -- HEADER -- | |
| st.title("🧠 TARS.help") | |
| st.markdown("### A minimal AI that listens, reflects, and replies.") | |
| st.markdown(f"🗓️ {datetime.now().strftime('%B %d, %Y')} | {len(st.session_state.history)//2} exchanges") | |
| # -- HIGH-RISK PHRASE FILTER -- | |
| TRIGGER_PHRASES = ["kill myself", "end it all", "suicide", "not worth living", "can't go on"] | |
| def is_high_risk(text): | |
| return any(phrase in text.lower() for phrase in TRIGGER_PHRASES) | |
| # -- INPUT -- | |
| user_input = st.text_input("How are you feeling today?", placeholder="Start typing...") | |
| # -- REPLY FUNCTION -- | |
| def generate_reply(user_input, context): | |
| prompt = f"""You are a calm, helpful AI assistant who supports users emotionally. Be kind and thoughtful in your reply. | |
| {context} | |
| User: {user_input} | |
| AI:""" | |
| response = generator(prompt, max_new_tokens=80, temperature=0.7)[0]['generated_text'] | |
| return response.split("AI:")[-1].strip() | |
| # -- CONVERSATION FLOW -- | |
| if user_input: | |
| context = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history[-4:]]) | |
| with st.spinner("TARS is reflecting..."): | |
| time.sleep(0.5) | |
| if is_high_risk(user_input): | |
| response = "I'm really sorry you're feeling this way. You're not alone — please talk to someone you trust or a mental health professional. 💙" | |
| else: | |
| response = generate_reply(user_input, context) | |
| timestamp = datetime.now().strftime("%H:%M") | |
| st.session_state.history.append(("🧍 You", user_input, timestamp)) | |
| st.session_state.history.append(("🤖 TARS", response, timestamp)) | |
| # -- DISPLAY HISTORY -- | |
| st.markdown("## 🗨️ Session") | |
| for speaker, msg, time in st.session_state.history: | |
| st.markdown(f"**{speaker} [{time}]:** {msg}") | |
| # -- SESSION SUMMARY -- | |
| if st.button("🧾 Generate Session Summary"): | |
| convo = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history]) | |
| prompt = f"""You are summarizing a thoughtful conversation between a user and an AI assistant. Write a kind, reflective note based on this interaction. | |
| Conversation: | |
| {convo} | |
| Summary:""" | |
| summary = generator(prompt, max_new_tokens=100, temperature=0.5)[0]['generated_text'] | |
| st.session_state.summary = summary.split("Summary:")[-1].strip() | |
| if st.session_state.summary: | |
| st.markdown("### 🧠 Session Note") | |
| st.markdown(st.session_state.summary) | |
| st.download_button("📥 Download Summary", st.session_state.summary, file_name="tars_session.txt") | |
| # -- FOOTER -- | |
| st.markdown("---") | |
| st.caption("TARS is not a therapist but a friend. Just a quiet assistant that reflects with you.") | |