Update app.py

app.py

@@ -1,63 +1,32 @@
 import streamlit as st
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-from sentence_transformers import SentenceTransformer
-import faiss
-import torch
 
-# Title of the Streamlit app
-st.title("Mental Health Chatbot")
-
-# Load a pre-trained sentence transformer model for embedding
-st.write("Loading models... Please wait.")
-sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
-
-# Load the RAG model, tokenizer, and retriever
-tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", use_dummy_dataset=True)
-rag_model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
+# Load the RAG model components
+@st.cache_resource
+def load_rag_model():
+    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+    retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", use_dummy_dataset=True)
+    rag_model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
+    return tokenizer, retriever, rag_model
 
-# Sample dialogues related to mental health (replace with actual dataset for production)
-sample_dialogues = [
-    "I'm feeling really down lately and don't know what to do.",
-    "I just lost my job, and I'm worried about the future.",
-    "I'm having trouble sleeping and feeling anxious all the time.",
-    "I've been feeling isolated and lonely.",
-    "I don't have the energy to do anything, and it's affecting my work."
-]
+tokenizer, retriever, rag_model = load_rag_model()
 
-# Embed the sample dialogues using the sentence transformer model
-embeddings = sentence_model.encode(sample_dialogues, convert_to_tensor=True)
-
-# Build FAISS index
-index = faiss.IndexFlatL2(embeddings.shape[1])
-index.add(embeddings.cpu().numpy())
+# Streamlit UI for Mental Health Chatbot
+st.title("Mental Health Chatbot")
+st.write("""
+This chatbot uses a pre-trained RAG model to provide responses to mental health-related queries. 
+Please note that this is an AI-based tool and is not a substitute for professional mental health support.
+""")
 
 # User input
-user_input = st.text_input("How are you feeling today?")
-
-# Define response generation function
-def generate_response(query):
-    # Embed the query using the sentence transformer
-    query_embedding = sentence_model.encode(query, convert_to_tensor=True).cpu().numpy()
-
-    # Search for the closest dialogue in the index
-    D, I = index.search(query_embedding, k=1)
-    
-    # Retrieve the closest dialogue
-    closest_dialogue = sample_dialogues[I[0][0]]
-
-    # Generate response using RAG model
-    inputs = tokenizer(closest_dialogue, return_tensors="pt")
-    outputs = rag_model.generate(**inputs)
-    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    
-    return response[0]
-
-# Generate a response when the user submits input
-if st.button("Talk to the Chatbot"):
-    if user_input:
-        with st.spinner('Generating response...'):
-            response = generate_response(user_input)
-            st.write(f"Chatbot: {response}")
+query = st.text_input("How can I help you today?")
+
+if st.button("Get Response"):
+    if query:
+        # Generate a response using the RAG model
+        inputs = tokenizer(query, return_tensors="pt")
+        outputs = rag_model.generate(**inputs)
+        response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        st.write(f"**Response:** {response[0]}")
     else:
-        st.write("Please enter something for the chatbot to respond to.")
+        st.write("Please enter a query to get a response.")