Spaces:
Sleeping
Sleeping
import streamlite as st | |
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer | |
# Load chatbot model | |
chatbot_model = "microsoft/DialoGPT-medium" | |
tokenizer = AutoTokenizer.from_pretrained(chatbot_model) | |
model = AutoModelForCausalLM.from_pretrained(chatbot_model) | |
# Load emotion detection model | |
emotion_pipeline = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base") | |
st.title("🧠 Mental Health Chatbot") | |
# Chat history | |
if "chat_history" not in st.session_state: | |
st.session_state.chat_history = [] | |
# User Input | |
user_input = st.text_input("You:", key="user_input") | |
if st.button("Send"): | |
if user_input: | |
# Generate chatbot response | |
input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt") | |
output = model.generate(input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id) | |
response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True) | |
# Detect emotion | |
emotion_result = emotion_pipeline(user_input) | |
emotion = emotion_result[0]["label"] | |
# Store chat history | |
st.session_state.chat_history.append(("You", user_input)) | |
st.session_state.chat_history.append(("Bot", response)) | |
# Display chat | |
for sender, msg in st.session_state.chat_history: | |
st.write(f"**{sender}:** {msg}") | |
# Display emotion | |
st.write(f"🧠 **Emotion Detected:** {emotion}") | |