Update rag_pipeline.py

rag_pipeline.py

@@ -1,16 +1,18 @@
+from datasets import load_dataset
+import pandas as pd
 from sentence_transformers import SentenceTransformer
 import faiss
-import numpy as np
-import pandas as pd
 from transformers import pipeline
 
 class RAGPipeline:
-    def __init__(self, dataset_path):
+    def __init__(self):
         self.embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
         self.generator = pipeline("text2text-generation", model="google/flan-t5-base")
-        self.data = pd.read_csv(dataset_path)
-        self.documents = self.data['context'].tolist()
-        self.questions = self.data['question'].tolist()
+
+        # Load dataset directly
+        ds = load_dataset("pubmed_qa", "pqa_labeled", split="train[:500]")
+        self.documents = ds["context"]
+        self.questions = ds["question"]
 
         self.index = self.build_faiss_index()
 
@@ -28,6 +30,6 @@ class RAGPipeline:
     def generate_answer(self, query):
         docs = self.retrieve(query)
         context = " ".join(docs)
-        prompt = f"Answer the following question using the provided context:\nContext: {context}\nQuestion: {query}"
+        prompt = f"Answer the following medical question using the context:\nContext: {context}\nQuestion: {query}"
         result = self.generator(prompt, max_length=200, do_sample=True)
         return result[0]['generated_text']