Update app.py

app.py

@@ -1,45 +1,22 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import pipeline
 
-# 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")
+# Load model from Hugging Face Hub
+model_name = "thrishala/mental_health_chatbot"
+nlp = pipeline("text-generation", model=model_name)
 
-# 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:"
+# Function to process user input and return chatbot's response
+def chatbot_response(user_input):
+    # Get response using the model pipeline
+    response = nlp(user_input, max_length=150, num_return_sequences=1)
+    return response[0]['generated_text']
 
-    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
-    )
+# Gradio interface
+iface = gr.Interface(fn=chatbot_response,
+                     inputs="text",
+                     outputs="text",
+                     title="Mental Health Chatbot",
+                     description="This chatbot provides empathetic responses to mental health-related queries. It aims to support users in a safe and compassionate manner.")
 
-    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()
+# Launch the app
+iface.launch()