| 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.from_dict(document) for document in documents] |
| return self.document_embedder.run(documents) |
|
|
| return {"error": "Wrong inputs"} |
|
|
