Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 from pathlib import Path
-from transformers import RagTokenForGeneration, RagTokenizer, DenseRetriever
+from transformers import RagTokenForGeneration, RagTokenizer
 import faiss
 from typing import List
 from pdfplumber import open as open_pdf
@@ -29,8 +29,19 @@ embeddings = rag_model.question_encoder(rag_tokenizer(text_chunks, padding=True,
 index = faiss.IndexFlatL2(embeddings.size(-1))
 index.add(embeddings.detach().numpy())
 
-# Create a custom retriever
-retriever = DenseRetriever(document_store=text_chunks, embedding=embeddings, similarities=index.search)
+# Custom retriever class
+class CustomRetriever:
+    def __init__(self, documents, embeddings, index):
+        self.documents = documents
+        self.embeddings = embeddings
+        self.index = index
+
+    def get_relevant_docs(self, query_embeddings, top_k=4):
+        scores, doc_indices = self.index.search(query_embeddings.detach().numpy(), top_k)
+        return [(self.documents[doc_idx], score) for doc_idx, score in zip(doc_indices[0], scores[0])]
+
+# Create a custom retriever instance
+retriever = CustomRetriever(text_chunks, embeddings, index)
 
 def respond(
     message,
@@ -52,8 +63,9 @@ def respond(
 
     # Retrieve relevant chunks using the custom retriever
     rag_input_ids = rag_tokenizer(message, return_tensors="pt").input_ids
-    rag_output = rag_model(rag_input_ids, retriever=retriever, return_retrieved_inputs=True)
-    retrieved_text = rag_output.retrieved_inputs
+    query_embeddings = rag_model.question_encoder(rag_input_ids)
+    relevant_docs = retriever.get_relevant_docs(query_embeddings)
+    retrieved_text = "\n".join([doc for doc, _ in relevant_docs])
 
     # Generate the response using the zephyr model
     for message in client.chat_completion(