syarulzaffi/date2format

1from typing import Dict, List, Any2from transformers import pipeline, AutoTokenizer3 4class EndpointHandler:5    def __init__(self, path=""):6        # Load the optimized model7        tokenizer = AutoTokenizer.from_pretrained(path)8        # Create inference pipeline for text classification9        self.pipeline = pipeline("text-classification", model=path, tokenizer=tokenizer)10 11    def __call__(self, data: str) -> List[List[Dict[str, float]]]:12        """13        Args:14            data (str): A raw string input for inference.15        Returns:16            A list containing the prediction results:17            A list of one list, e.g., [[{"label": "LABEL", "score": 0.99}]]18        """19        # Pass the data as `text` directly20        inputs = data.pop("inputs", data)21        prediction = self.pipeline(inputs)22 23        # Return the prediction result24        return prediction