Update search.py

search.py

@@ -4,13 +4,21 @@ from pdfminer.high_level import extract_text
 from typing import List
 import pandas as pd
 import re
+from datasets import load_dataset
 
 # Initialize RAG components
 rag_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
-rag_retriever = RagRetriever.from_pretrained("facebook/rag-token-base", trust_remote_code=True)
 rag_token_for_generation = RagTokenForGeneration.from_pretrained("facebook/rag-token-base")
 rag_config = RagConfig.from_pretrained("facebook/rag-token-base")
 
+# Download and prepare the wiki_dpr dataset
+dpr_dataset = load_dataset("wiki_dpr")
+passages = dpr_dataset["train"]["passage"]
+titles = dpr_dataset["train"]["title"]
+
+# Initialize the RagRetriever
+rag_retriever = RagRetriever(passages=passages, titles=titles, config=rag_config)
+
 # Dataclass for paragraph
 @dataclass
 class Paragraph:
@@ -114,13 +122,10 @@ def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, retrie
     
     context = "\n\n".join(most_relevant_paragraphs)
     
-    # Retrieve documents relevant to the question
-    documents = retriever.retrieve(question)
-    
     # Generate answer using RAG
-    inputs = tokenizer(question, context, return_tensors="pt", max_length=512, truncation=True)
-    outputs = generator.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_length=200, num_return_sequences=1)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = rag_tokenizer(question, context, return_tensors="pt", max_length=512, truncation=True)
+    outputs = rag_token_for_generation.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_length=200, num_return_sequences=1)
+    answer = rag_tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     references = extract_page_and_clause_references(context)
     answer = refine_answer_based_on_question(question, answer) + " " + references