Update app.py

app.py

@@ -1,8 +1,10 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 
-# Load the model and tokenizer
-chatbot = pipeline("text2text-generation", model="google/flan-t5-large")
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it", torch_dtype=torch.bfloat16)
 
 # Set up the Streamlit page configuration
 st.set_page_config(page_title="AI Companion Chatbot", layout="centered")
@@ -19,25 +21,34 @@ providing a safe and empathetic space for you to express your feelings.
 # Create a text input box for user input
 user_input = st.text_area("How are you feeling today?", "")
 
+# Define the function to generate the response
+def generate_response(user_input):
+    prompt = f"""
+    You are a therapist with a strong focus on providing practical, actionable advice.
+    Rules:
+    1. Respond in a supportive, empathetic, and non-judgmental manner to the following statement.
+    2. Offer at least 3 **specific** strategies or coping techniques that the user can try immediately to manage or alleviate their anxiety. 
+    These could include emotional regulation techniques (like grounding exercises, breathing techniques), 
+    self-care practices (like self-compassion or taking breaks), or mindset shifts (like reframing negative thoughts or focusing on what can be controlled).
+    3. Be very descriptive. Use bullet points to clearly state actionable steps.
+    4. Do not use "I" or reference the first person perspective.
+    
+    Base your response on how the user is feeling: {user_input}
+    """
+
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+    # Generate the output
+    outputs = model.generate(**inputs, max_length=350, num_return_sequences=1)
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
 # Display chatbot's response
 if st.button("Send"):
     if user_input:  # Check if the user has provided input
-        # Construct the message to instruct the model to behave like a therapist
-        prompt = f"""
-        You are a therapist with a strong focus on providing practical, actionable advice.
-        Rules:
-        1. Respond in a supportive, empathetic, and non-judgmental manner to the following statement.
-        2. Offer at least 3 **specific** strategies or coping techniques that the user can try immediately to manage or alleviate their anxiety. 
-        These could include emotional regulation techniques (like grounding exercises, breathing techniques), 
-        self-care practices (like self-compassion or taking breaks), or mindset shifts (like reframing negative thoughts or focusing on what can be controlled).
-        3. Be very descriptive. Use bullet points to clearly state actionable steps.
-        4. Do not use "I" or reference the first person perspective.
-        
-        Base your response on how the user is feeling: {user_input}
-        """
-
         # Get the response from the model
-        response = chatbot(prompt, max_length=350, num_return_sequences=1)[0]['generated_text']
+        response = generate_response(user_input)
 
         # Show the response
         st.text_area("AI Companion Response:", response, height=200)