Spaces:
Sleeping
Sleeping
File size: 1,535 Bytes
db6eba4 bcee0a0 db6eba4 bcee0a0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load the model and tokenizer
model_name = "givyboy/TinyLlama-1.1B-Chat-v1.0-mental-health-conversational"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
# Function to generate response
def chat_with_bot(message, history):
history = history or []
chat_history = ""
for user_msg, bot_msg in history:
chat_history += f"User: {user_msg}\nBot: {bot_msg}\n"
chat_history += f"User: {message}\nBot:"
inputs = tokenizer(chat_history, return_tensors="pt", return_attention_mask=False).to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=150,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
bot_reply = response.split("Bot:")[-1].strip()
history.append((message, bot_reply))
return bot_reply, history
# UI
chat_interface = gr.ChatInterface(
fn=chat_with_bot,
title="🧠 SerenityAI Mental Health Chatbot",
description="Welcome! I'm here to help you talk through your thoughts and feelings. 💙",
theme="soft",
examples=["I'm feeling anxious lately.", "I had a great day!", "I feel really overwhelmed."],
retry_btn="🔁 Try Again",
clear_btn="🧹 Clear Chat",
)
chat_interface.launch()
|