k4tel/bert-geolocation-prediction

1from typing import Dict, List, Any2import torch 3from transformers import pipeline4 5class EndpointHandler():6    def __init__(self, path=""):7        self.pipeline = pipeline("geo-regression", model=path)8 9    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:10        """11         data args:12              inputs (:obj: `str`)13              date (:obj: `str`)14        Return:15              A :obj:`list` | `dict`: will be serialized and returned16        """17        # get inputs18        inputs = data.pop("inputs", data)19        # get additional date field20        date = data.pop("date", None)21 22        # check if date exists and if it is a holiday23 24        # run normal prediction25        prediction = self.pipeline(inputs)26        return prediction27 28