Update search.py

search.py

@@ -1,16 +1,12 @@
-from transformers import RagTokenizer, RagTokenForGeneration, AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import RagTokenizer, RagTokenForGeneration, pipeline
+import pandas as pd
 from pdfminer.high_level import extract_text
 from docx import Document
 from dataclasses import dataclass
-import pandas as pd
 
-# Initialize RAG
-rag_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
-rag_model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")
-
-# Initialize Phi-2
-phi_tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-phi_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
+# RAG setup
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")
 
 @dataclass
 class Paragraph:
@@ -18,41 +14,19 @@ class Paragraph:
     paragraph_num: int
     content: str
 
-def read_pdf_pdfminer(file_path) -> list[Paragraph]:
+def read_pdf_pdfminer(file_path) -> list:
     text = extract_text(file_path).replace('\n', ' ').strip()
-    paragraphs = text.split(". ")
-    return [Paragraph(0, i, para) for i, para in enumerate(paragraphs, 1)]
+    paragraphs = [Paragraph(0, i, para) for i, para in enumerate(text.split('. '), start=1)]
+    return paragraphs
 
-def read_docx(file) -> list[Paragraph]:
+def read_docx(file) -> list:
     doc = Document(file)
-    return [Paragraph(1, i, para.text.strip()) for i, para in enumerate(doc.paragraphs, 1) if para.text.strip()]
-
-def generate_context_with_rag(question: str, documents: List[str]) -> str:
-    combined_text = " ".join(documents)
-    if not combined_text.strip():  # Ensure combined_text is not empty
-        return "No context available."
-
-    inputs = rag_tokenizer(question + " " + combined_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    
-    # Ensure inputs are correctly prepared
-    if "input_ids" not in inputs or "attention_mask" not in inputs:
-        return "Invalid input for model."
-    
-    output_ids = rag_model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
-    context = rag_tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return context
-
-
-def generate_answer_with_phi(question: str, context: str) -> str:
-    enhanced_question = f"Question: {question}\nContext: {context}\nAnswer:"
-    inputs = phi_tokenizer.encode(enhanced_question, return_tensors="pt", max_length=512, truncation=True)
-    outputs = phi_model.generate(inputs, max_length=600)
-    answer = phi_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer
+    return [Paragraph(1, i, para.text.strip()) for i, para in enumerate(doc.paragraphs, start=1) if para.text.strip()]
 
 def answer_question(question: str, documents_df: pd.DataFrame) -> str:
-    # Assuming documents_df contains the text from uploaded files
-    combined_text = " ".join(documents_df['content'].tolist())
-    context = generate_context_with_rag(combined_text + " " + question)
-    answer = generate_answer_with_phi(question, context)
+    document_texts = " ".join(documents_df['content'].tolist())
+    context = f"{question} {document_texts}"
+    inputs = tokenizer(context, return_tensors="pt", truncation=True, max_length=512, padding="max_length")
+    output_ids = model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
+    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return answer