Update handler.py

handler.py

@@ -1,43 +1,23 @@
-from typing import Any, Dict, List
+from typing import Any, Dict
 
-from haystack.schema import Document
-from fastrag.rankers import QuantizedBiEncoderRanker
+from haystack import Document
+
+from fastrag.rankers import IPEXBiEncoderSimilarityRanker
 
 
 class EndpointHandler:
     def __init__(self, path=""):
         model_id = "Intel/bge-large-en-v1.5-rag-int8-static"
-        self.ranker = QuantizedBiEncoderRanker(model_name_or_path=model_id)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        query = data.get("query", None)
-        queries = data.get("queries", None)
-        documents = data.get("documents", None)
-        batch_size = data.get("batch_size", None)
+        self.ranker = IPEXBiEncoderSimilarityRanker(model=model_id)
+
+        self.ranker.warm_up()
+
+    def __call__(self, data: Dict[str, Any]):
+        query = data.get("query")
+        documents = data.get("documents")
         top_k = data.get("top_k", None)
 
-        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
-            )
-
-        else:
-            raise ValueError("Expected either query or queries")
+        documents = [Document.from_dict(doc) for doc in documents]
+
+        return self.ranker.run(query, documents, top_k)