| from typing import Any, Dict, List | |
| from haystack import Document | |
| from fastrag.embedders import ( | |
| IPEXSentenceTransformersDocumentEmbedder, | |
| IPEXSentenceTransformersTextEmbedder, | |
| ) | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| model_id = "Intel/bge-small-en-v1.5-rag-int8-static" | |
| self.query_embedder = IPEXSentenceTransformersTextEmbedder(model=model_id) | |
| self.document_embedder = IPEXSentenceTransformersDocumentEmbedder(model=model_id) | |
| self.query_embedder.warm_up() | |
| self.document_embedder.warm_up() | |
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| text = data.get("text", None) | |
| if text is not None: | |
| assert isinstance(text, str), "Expected text to be a string" | |
| return self.query_embedder.run(text) | |
| documents = data.get("documents", None) | |
| if documents is not None: | |
| assert isinstance(documents, list), "Expected documents to be a list" | |
| assert all( | |
| isinstance(document, dict) for document in documents | |
| ), "Expected each document in documents to be a dictionary" | |
| documents = [Document(**document) for document in documents] | |
| return self.document_embedder.run(documents) | |
| raise ValueError("Expected either text or documents") | |