| from typing import Any, Dict, List | |
| 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") | |
| documents = data.get("documents") | |
| 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" | |
| return self.ranker.predict(query=query, documents=documents, top_k=top_k) | |