app.py

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()