Update app.py

app.py

@@ -2,8 +2,6 @@ import streamlit as st
 import fitz  # PyMuPDF
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 import numpy as np
-import faiss
-import torch
 
 # Load the RAG model components
 tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
@@ -24,7 +22,7 @@ def answer_question(question, pdf_text):
     inputs = tokenizer(question, return_tensors="pt")
     
     # Retrieve documents based on the PDF text
-    doc_embeds = retriever.get_document_embeddings(pdf_text)
+    doc_embeds = retriever.get_document_embeddings([pdf_text])  # Wrap pdf_text in a list
     retriever.set_retriever_doc_embeddings(doc_embeds)
     
     # Get the top k documents for the question
@@ -63,4 +61,3 @@ if pdf_file is not None:
 
 
 
-