Update app.py

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-from transformers import pipeline
+from transformers import RagTokenizer, RagTokenForGeneration
 from typing import List, Dict, Tuple
 import re
 import os
+import torch
 
-# Set up the retriever pipeline
-retriever = pipeline('retrieval', model='facebook/rag-token-nq')
+# Load the RAG model and tokenizer
+rag_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+rag_model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")
 
 # Load your PDF document
 pdf_path = "apexcustoms.pdf"
@@ -44,10 +46,12 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    # Retrieve relevant context from the PDF
-    retrieval_output = retriever(message, corpus, top_k=3)
-    retrieved_contexts = [passage['text'] for passage in retrieval_output['retrieved_passages']]
-    context = ' '.join(retrieved_contexts)
+    # Tokenize the input and retrieve relevant context from the PDF
+    inputs = rag_tokenizer(message, return_tensors="pt")
+    inputs.update({"corpus": corpus})
+    input_ids = inputs.pop("input_ids")
+    output_ids = rag_model.generate(**inputs, max_length=max_tokens, temperature=temperature, top_p=top_p, num_beams=2)
+    retrieved_context = rag_tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
 
     response = ""
 
@@ -57,7 +61,7 @@ def respond(
         stream=True,
         temperature=temperature,
         top_p=top_p,
-        context=context,  # Include the retrieved context
+        context=retrieved_context,  # Include the retrieved context
     ):
         token = message.choices[0].delta.content