Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 import fitz  # PyMuPDF
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-import numpy as np
+import torch
 
 # Load the RAG model components
 tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
@@ -18,24 +18,19 @@ def extract_text_from_pdf(pdf_file):
 
 # Function to handle question answering
 def answer_question(question, pdf_text):
-    # Tokenize the question
-    inputs = tokenizer(question, return_tensors="pt")
-    
-    # Retrieve documents based on the 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
-    k = 5
-    retrieved_docs = retriever(question, n_docs=k)
-    
     # Prepare the context for the model
-    context = retrieved_docs["document_texts"]
-    context = " ".join(context)
+    inputs = tokenizer([question], return_tensors="pt")
     
+    # Tokenize PDF text
+    pdf_inputs = tokenizer([pdf_text], return_tensors="pt")
+
     # Generate the answer
-    input_dict = tokenizer.prepare_seq2seq_batch(question, context, return_tensors="pt")
-    outputs = model.generate(**input_dict)
+    with torch.no_grad():
+        outputs = model.generate(input_ids=inputs['input_ids'], 
+                                 attention_mask=inputs['attention_mask'], 
+                                 context_input_ids=pdf_inputs['input_ids'], 
+                                 context_attention_mask=pdf_inputs['attention_mask'])
+    
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return answer
 
@@ -55,9 +50,12 @@ if pdf_file is not None:
     
     if question:
         with st.spinner("Finding answer..."):
-            answer = answer_question(question, pdf_text)
-            st.write("### Answer:")
-            st.write(answer)
+            try:
+                answer = answer_question(question, pdf_text)
+                st.write("### Answer:")
+                st.write(answer)
+            except Exception as e:
+                st.error(f"Error occurred: {str(e)}")