Update app.py

app.py

@@ -1,64 +1,56 @@
 import gradio as gr
-from datasets import load_dataset
-import random
-import groq  # Assuming you are using the Groq library
-
-# Load dataset
-ds = load_dataset("alexandreteles/mental-health-conversational-data")
-
-# Extract columns
-context = ds["train"]["Context"]
-knowledge = ds["train"]["Knowledge"]
-response = ds["train"]["Response"]
-
-# Directly input the Groq API key (replace with your actual API key)
-api_key = "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"
-
-# Initialize Groq API client with the API key
-client = groq.Client(api_key=api_key)
-
-# Define a function to simulate a conversation
-def chatbot(user_input):
-    if not user_input.strip():
-        return "Please enter a question or concern to receive guidance."
-
-    try:
-        # Try to call Groq API to generate a response
-        brief_response = client.predict(user_input)  # Ensure that 'predict' or correct method is available
-    except Exception as e:
-        brief_response = None  # If Groq API fails, set brief_response to None
-
-    if not brief_response:
-        # If Groq API does not return a response, fall back to dataset
-        idx = random.randint(0, len(context)-1)
-        reply = response[idx]
-        context_text = context[idx]
-        knowledge_text = knowledge[idx]
-        
-        # Combine dataset info with fallback response
-        complete_response = (
-            f"**Contextual Information**\n{context_text}\n\n"
-            f"**Knowledge Base**\n{knowledge_text}\n\n"
-            f"**Fallback Response**\n{reply}"
-        )
-    else:
-        # If Groq API returns a response, use that as the final output
-        complete_response = f"**Personalized Response**\n{brief_response}"
-
-    return complete_response
-
-# Enhance the interface
-interface = gr.Interface(
-    fn=chatbot,
-    inputs=gr.Textbox(label="Ask your question:", placeholder="How are you feeling today?"),
-    outputs=gr.Markdown(label="Psychologist Assistant Response"),
-    title="Virtual Psychologist Assistant",
-    description=(
-        "This is a supportive assistant designed to provide compassionate guidance "
-        "for mental well-being. Type your thoughts or questions for tailored advice and insights."
-    ),
-    theme="huggingface",  # Optional: apply a theme if available
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
+import torch
+import os
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+HUGGING_FACE_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
+
+# Load the sentiment analysis model
+sentiment_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+sentiment_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+
+# Load the LLaMA-1B model for text generation, using the token from the environment variable
+llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", use_auth_token=HUGGING_FACE_TOKEN)
+llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", use_auth_token=HUGGING_FACE_TOKEN)
+
+# Function for sentiment analysis
+def analyze_sentiment(text):
+    inputs = sentiment_tokenizer(text, return_tensors="pt")
+    outputs = sentiment_model(**inputs)
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    sentiment = "positive" if torch.argmax(probs) == 1 else "negative"
+    confidence = probs.max().item()
+    return sentiment, confidence
+
+# Function to generate a supportive response
+def generate_response(sentiment, text):
+    prompt = f"The user feels {sentiment}. Respond with supportive advice based on: {text}"
+    inputs = llama_tokenizer(prompt, return_tensors="pt")
+    response = llama_model.generate(**inputs, max_length=150)
+    return llama_tokenizer.decode(response[0], skip_special_tokens=True)
+
+# Combine both functions for the personal psychologist
+def personal_psychologist(text):
+    sentiment, confidence = analyze_sentiment(text)
+    response = generate_response(sentiment, text)
+    return f"Sentiment: {sentiment} (Confidence: {confidence:.2f})\n\nResponse: {response}"
+
+# Set up Gradio interface
+iface = gr.Interface(
+    fn=personal_psychologist,
+    inputs="text",
+    outputs="text",
+    title="Personal Psychologist",
+    description="A supportive AI that assesses your mood and provides comforting advice based on your input.",
+    examples=[
+        ["I'm feeling very anxious and stressed about my exams."],
+        ["I had a great day with my friends!"],
+        ["I feel like I'm not good enough and everything is going wrong."]
+    ]
 )
 
-# Launch the app
-interface.launch()
+# Launch Gradio app
+iface.launch()