Update handler.py

handler.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, List
 
+from haystack.schema import Document
 from fastrag.rankers import QuantizedBiEncoderRanker
 
 
@@ -9,13 +10,34 @@ class EndpointHandler:
         self.ranker = QuantizedBiEncoderRanker(model_name_or_path=model_id)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        query = data.get("query")
-        documents = data.get("documents")
-
+        query = data.get("query", None)
+        queries = data.get("queries", None)
+        documents = data.get("documents", None)
+        batch_size = data.get("batch_size", None)
         top_k = data.get("top_k", None)
 
-        assert isinstance(query, str), "Expected query to be a string"
-        assert isinstance(documents, list), "Expected documents to be a list"
-        assert all(isinstance(document, str) for document in documents), "Expected documents to be a list of strings"
+        if query is not None:
+            assert isinstance(query, str), "Expected query to be a string"
+            assert isinstance(documents, list), "Expected documents to be a list"
+            assert all(
+                isinstance(d, dict) for d in documents
+            ), "Expected each document in documents to be a dictionary"
+            documents = [Document.from_dict(d) for d in documents]
+            return self.ranker.predict(query=query, documents=documents, top_k=top_k)
+
+        elif queries is not None:
+            assert isinstance(queries, list), "Expected queries to be a list"
+            assert all(
+                isinstance(query, str) for query in queries
+            ), "Expected each query in queries to be a string"
+            assert isinstance(documents, list), "Expected documents to be a list"
+            assert all(
+                all(isinstance(d, dict) for d in doc) for doc in documents
+            ), "Expected each document in list of documents to be a dictionary"
+            documents = [Document.from_dict(d) for d in documents]
+            return self.ranker.predict_batch(
+                queries=queries, documents=documents, batch_size=batch_size, top_k=top_k
+            )
 
-        return self.ranker.predict(query=query, documents=documents, top_k=top_k)
+        else:
+            raise ValueError("Expected either query or queries")