add custom handler

handler.py

@@ -0,0 +1,63 @@
+import logging
+from typing import List, Dict
+
+from pydantic import BaseModel
+from sentence_transformers import CrossEncoder
+
+import config
+
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.DEBUG)
+console_handler.setFormatter(formatter)
+logger = logging.getLogger()
+logger.addHandler(console_handler)
+logger.setLevel(logging.DEBUG)
+
+
+class DocsToRerank(BaseModel):
+    query: str
+    documents: List[str]
+
+
+class Payload(BaseModel):
+    inputs: DocsToRerank
+    parameters: Dict[str, str] = None
+
+
+class Rank(BaseModel):
+    corpus_id: str
+    score: float
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        crossencoder_model = config.CROSSENCODER_MODEL
+        self.top_k = config.CROSSENCODER_TOP_K
+        self.max_length = config.CROSSENCODER_MAX_LENGTH
+        self.max_workers = config.CROSSENCODER_MAX_WORKERS
+
+        self.model = CrossEncoder(
+            crossencoder_model,
+            max_length=self.max_length
+        )
+        logger.info(
+            f'Loaded {path}:top_k={self.top_k}:length={self.max_length}:workers={self.max_workers}')
+
+    def __call__(self, data: Payload) -> List[Rank]:
+        data = data.inputs
+
+        logger.info(f'Received new docs to rerank: {data}')
+        ranks = self.model.rank(
+            data.query,
+            data.documents,
+            top_k=self.top_k,
+            num_workers=self.max_workers
+        )
+
+        ranks = [Rank(
+            corpus_id=r['corpus_id'],
+            score=r['score']) for r in ranks
+        ]
+        logger.info(f'Computed ranks for each document: {ranks}')
+        return ranks

requirements.txt

@@ -0,0 +1,3 @@
+sentence-transformers~=3.0.0
+pydantic
+pytest

test_endpoint_handler.py

@@ -0,0 +1,19 @@
+import pytest
+
+from handler import EndpointHandler, Payload
+
+
+@pytest.mark.parametrize('payload', [
+    {
+        'inputs': {
+            'query': 'I like apples',
+            'documents': [
+                'But I only eat the red ones',
+                'Asia is a continent',
+                'You can add them to smoothies'
+            ]
+        }
+    }
+])
+def test_sanity(payload):
+    assert EndpointHandler(path=".")(Payload.model_validate(payload))