| from typing import Dict, List, Any | |
| from comet import load_from_checkpoint | |
| import os | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| os.listdir(os.cwd()) | |
| self.model = load_from_checkpoint("model.ckpt") | |
| def __call__(self, data: Dict[str, Any]) -> List[Any]: | |
| """ | |
| data args: | |
| inputs (:obj: `dict[str, str]`) | |
| Return: | |
| A :obj: `dict`: will be serialized and returned | |
| """ | |
| inputs = data.pop("inputs") | |
| batch_size = inputs.pop("batch_size") | |
| workers = inputs.pop("workers") | |
| data = inputs.pop("data") | |
| model_output = self.model.predict(data, batch_size=batch_size, num_workers=workers, gpus=0) | |
| scores = model_output["scores"] | |
| return scores | |