| from typing import Any, Dict, List | |
| from haystack.schema import Document | |
| from fastrag.rankers import QuantizedBiEncoderRanker | |
| 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) | |
| 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") | |