RAG fix

search.py

@@ -27,12 +27,22 @@ def read_docx(file) -> list[Paragraph]:
     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) -> str:
-    inputs = rag_tokenizer(question, return_tensors="pt")
-    output_ids = rag_model.generate(**inputs)
+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)