Update app.py

app.py

@@ -1,68 +1,32 @@
-#To add temperature control, increase token lengths, and restrictions to provide suggestions to visit nearby therapists, you can modify the chatbot function as follows:
-
-#app.py (Streamlit version with temperature control and restrictions)
-
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-# Load pre-trained model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("bert-base-uncased")
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-
-def chatbot(input_message, temperature=0.7, max_length=50):
-    # System roles
-    if input_message.lower() == "hello":
-        return "Welcome to the Mental Health Chatbot! I'm here to listen and help."
-    elif input_message.lower() == "help":
-        return "I'm here to provide support and resources for mental health. What's on your mind?"
-    # Assistant roles
-    elif input_message.lower().startswith("what is"):
-        return "I can provide information on various mental health topics. Please ask a specific question."
-    elif input_message.lower().startswith("how to"):
-        return "I can offer suggestions and resources for managing mental health. Please ask a specific question."
-    # Therapist suggestion
-    elif input_message.lower().startswith("i need help"):
-        return "Consider visiting a nearby therapist for professional support. You can search online for therapists in your area."
-    # Default response
-    else:
-        input_ids = tokenizer.encode(input_message, return_tensors="pt")
-        outputs = model.generate(
-            input_ids=input_ids,
-            max_length=max_length,
-            num_return_sequences=1,
-            temperature=temperature,
-            top_k=50,
-            top_p=0.95,
-            repetition_penalty=1.2,
-            no_repeat_ngram_size=3,
-        )
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Restrict responses to mental health topics
-        if "mental health" not in response.lower():
-            response = "I'm here to provide support and resources for mental health. Please ask a specific question."
-        return response
+from transformers import AutoModelWithLMHead, AutoTokenizer
+
+# Load pre-trained T5 base model and tokenizer
+model = AutoModelWithLMHead.from_pretrained("t5-base")
+tokenizer = AutoTokenizer.from_pretrained("t5-base")
+
+def chatbot(input_message):
+    input_ids = tokenizer.encode(f"generate text: {input_message}", return_tensors="pt")
+    outputs = model.generate(
+        input_ids=input_ids,
+        max_length=50,
+        num_return_sequences=1,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95,
+        repetition_penalty=1.2,
+        no_repeat_ngram_size=3,
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 def main():
     st.title("Mental Health Chatbot")
     input_message = st.text_input("You:")
-    temperature = st.slider("Temperature", 0.1, 1.0, 0.7)
-    max_length = st.slider("Max Length", 50, 100, 50)
     if st.button("Send"):
-        response = chatbot(input_message, temperature, max_length)
+        response = chatbot(input_message)
         st.text_area("Chatbot:", value=response, height=100)
 
 if __name__ == "__main__":
     main()
 
-# Run the app with: streamlit run app.py
-'''
-Here's what's changed:
-
-- I added two sliders to control the temperature and max length of the responses.
-- I added a conditional statement to suggest visiting a nearby therapist if the user inputs "I need help".
-- I restricted the responses to mental health topics by checking if the response contains "mental health".
-
-You can adjust the temperature and max length sliders to control the creativity and length of the responses. The temperature parameter controls how likely the model is to generate novel responses, while the max length parameter controls the maximum length of the responses.
-
-Note that you may need to adjust the model and tokenizer to better fit your specific use case. Additionally, you may want to consider adding more restrictions and guidelines to ensure the chatbot provides appropriate and safe responses.'''