instructor-base / handler.py
t12e's picture
updated handler.py
41c2439
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,
"instruction": instruction,
"document": document
}
]