| from typing import Any, Dict, List |
|
|
| from haystack.schema import Document |
| from fastrag.retrievers import QuantizedBiEncoderRetriever |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| model_id = "Intel/bge-small-en-v1.5-rag-int8-static" |
| self.retriever = QuantizedBiEncoderRetriever(embedding_model=model_id) |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
| queries = data.get("queries", None) |
| documents = data.get("documents", None) |
|
|
| if 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" |
|
|
| return self.retriever.embed_queries(queries=queries) |
| elif 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.retriever.embed_documents(documents=documents) |
| else: |
| raise ValueError("Expected either queries or documents") |
|
|
