Update app.py

app.py

@@ -1,41 +1,40 @@
 import streamlit as st
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 from datasets import load_dataset
-import torch
 
-# Load the dataset
+# Load dataset (pubmed_qa) and tokenizer
 dataset = load_dataset("pubmed_qa", split="test")
-
-# Initialize RAG components
 tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
-retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="default", use_dummy_dataset=True)
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
-
-# Function to get the answer to a medical query
-def get_medical_answer(query):
-    # Encode the query to retrieve relevant documents
-    inputs = tokenizer(query, return_tensors="pt")
-    input_ids = inputs["input_ids"]
-    
-    # Retrieve relevant documents
-    docs = retriever(input_ids=input_ids, return_tensors="pt")
-    
-    # Generate the answer from the model
-    generated_ids = model.generate(input_ids=input_ids, context_input_ids=docs["context_input_ids"], 
-                                   context_attention_mask=docs["context_attention_mask"])
-    
-    # Decode the generated answer
-    generated_answer = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    return generated_answer
-
-# Streamlit UI
-st.title("Medical QA Assistant")
-st.write("Ask any medical question, and I will answer it based on PubMed papers!")
-
-# Input text box for queries
-query = st.text_input("Enter your medical question:")
-
-if query:
-    with st.spinner("Searching for the answer..."):
-        answer = get_medical_answer(query)
+retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="compressed", passages_path="./path_to_dataset")
+
+# Initialize the RAG model
+model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
+
+# Define Streamlit app
+st.title('Medical QA Assistant')
+
+st.markdown("This app uses a RAG model to answer medical queries based on the PubMed QA dataset.")
+
+# User input for query
+user_query = st.text_input("Ask a medical question:")
+
+if user_query:
+    # Tokenize input question and retrieve related documents
+    inputs = tokenizer(user_query, return_tensors="pt")
+    input_ids = inputs['input_ids']
+    question_encoder_outputs = model.question_encoder(input_ids)
+
+    # Use the retriever to get context
+    retrieved_docs = retriever.retrieve(input_ids)
+
+    # Generate an answer based on the context
+    generated_ids = model.generate(input_ids, context_input_ids=retrieved_docs)
+    answer = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+    # Show the answer
     st.write(f"Answer: {answer}")
+
+    # Display the most relevant documents
+    st.subheader("Relevant Documents:")
+    for doc in retrieved_docs:
+        st.write(doc['text'][:300] + '...')  # Display first 300 characters of each doc