Update app.py

app.py

@@ -24,34 +24,48 @@ docs = prepare_docs(dataset)
 
 # Custom Retriever that searches in the dataset
 class HoroscopeRetriever(RagRetriever):
-    def __init__(self, docs):
+    def __init__(self, docs, tokenizer):
         self.docs = docs
+        self.tokenizer = tokenizer
 
     def retrieve(self, question_hidden_states, n_docs=1):
         # Convert the question_hidden_states to a text string
-        question_hidden_states = question_hidden_states[0]
+        question = question_hidden_states[0]
         
-        if isinstance(question_hidden_states, np.ndarray):
-            if question_hidden_states.size == 1:
-                question = question_hidden_states.item()  # Convert single-element array to scalar
+        if isinstance(question, np.ndarray):
+            if question.size == 1:
+                question = question.item()  # Convert single-element array to scalar
             else:
-                question = str(question_hidden_states[0])  # Take the first element of the array
+                question = str(question[0])  # Take the first element of the array
         else:
-            question = str(question_hidden_states)
+            question = str(question)
         
         question = question.lower()
         
-        # Simple retrieval logic: return the most relevant document based on the question
+        # Simple retrieval logic: find the most relevant document based on the question
+        best_match = None
         for doc in self.docs:
             if question in doc["question"].lower():
-                return [doc["answer"]]
-        return ["Sorry, I couldn't find a relevant horoscope."]
+                best_match = doc
+                break
+        
+        if best_match:
+            # Fake embedding as RAG expects this (In a real case, compute embeddings)
+            retrieved_doc_embeds = torch.zeros((1, 1, 768))  # Example tensor
+            doc_ids = ["0"]  # Example document ID
+            docs = [best_match["answer"]]
+        else:
+            retrieved_doc_embeds = torch.zeros((1, 1, 768))  # Example tensor
+            doc_ids = ["0"]  # Example document ID
+            docs = ["Sorry, I couldn't find a relevant horoscope."]
+        
+        return retrieved_doc_embeds, doc_ids, docs
 
 # Initialize the custom retriever with the dataset
-retriever = HoroscopeRetriever(docs)
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
+retriever = HoroscopeRetriever(docs, tokenizer)
 
 # Initialize RAG components
-tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
 model = RagTokenForGeneration.from_pretrained("facebook/rag-token-base", retriever=retriever)
 
 # Define the chatbot function