Update app.py

app.py

@@ -1,32 +1,27 @@
 import streamlit as st
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from transformers import RagRetriever, RagTokenizer, RagTokenForGeneration
 
-# Load the RAG model components
+# Load the RAG model
 @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)
+    retriever = RagRetriever.from_pretrained("facebook/rag-token-nq")
+    tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+    rag_model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
     return tokenizer, retriever, rag_model
 
-tokenizer, retriever, rag_model = load_rag_model()
+# Set up the Streamlit interface
+def main():
+    st.title("Mental Health Chatbot")
 
-# 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.
-""")
+    tokenizer, retriever, rag_model = load_rag_model()
 
-# User input
-query = st.text_input("How can I help you today?")
+    user_input = st.text_input("Ask me something about mental health:")
+    
+    if user_input:
+        input_ids = tokenizer(user_input, return_tensors="pt").input_ids
+        output = rag_model.generate(input_ids)
+        response = tokenizer.decode(output[0], skip_special_tokens=True)
+        st.write(f"Response: {response}")
 
-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 a query to get a response.")
+if __name__ == "__main__":
+    main()