instructor-base / handler.py
t12e's picture
add custom handler again
98c2017
raw
history blame
457 Bytes
from typing import Dict, List, Any
from InstructorEmbedding import INSTRUCTOR
class EndpointHandler():
def __init__(self, path=""):
self.model = INSTRUCTOR(path)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
instruction = data.get("instruction", "")
document = data.get("document", "")
embedding = self.model.encode([[instruction, document]]).flatten()
return [{"embedding": embedding}]