Update app.py

app.py

@@ -1,22 +1,32 @@
-import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Load model from Hugging Face Hub
-model_name = "thrishala/mental_health_chatbot"
-nlp = pipeline("text-generation", model=model_name)
+# Load the model and tokenizer
+model_name = "tanusrich/Mental_Health_Chatbot"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# 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']
+# Move the model to the appropriate device (CPU or GPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# 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.")
+# Generate a response
+def generate_response(user_input):
+    inputs = tokenizer(user_input, return_tensors="pt").to(device)
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.9,
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
 
-# Launch the app
-iface.launch()
+# Example interaction
+user_input = "I'm feeling lonely and anxious. What can I do?"
+response = generate_response(user_input)
+print("Chatbot: ", response)