Update app.py

app.py

@@ -1,23 +1,41 @@
 import streamlit as st
-from rag_pipeline import RAGPipeline
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from datasets import load_dataset
+import torch
 
-st.set_page_config(page_title="Medical QA Assistant", page_icon="🩺")
-st.title("🩺 Medical QA Assistant")
-st.markdown("Ask any medical question and get evidence-based answers from PubMed.")
+# Load the dataset
+dataset = load_dataset("pubmed_qa", split="test")
 
-@st.cache_resource
-def load_rag():
-    return rag = RAGPipeline()
+# 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
 
-rag = load_rag()
+# Streamlit UI
+st.title("Medical QA Assistant")
+st.write("Ask any medical question, and I will answer it based on PubMed papers!")
 
-question = st.text_input("Enter your medical question:")
+# Input text box for queries
+query = st.text_input("Enter your medical question:")
 
-if st.button("Get Answer"):
-    if question:
-        with st.spinner("Searching PubMed and generating answer..."):
-            answer = rag.generate_answer(question)
-            st.success(answer)
-    else:
-        st.warning("Please enter a question.")
+if query:
+    with st.spinner("Searching for the answer..."):
+        answer = get_medical_answer(query)
+    st.write(f"Answer: {answer}")