from typing import List, Optional, Union

from infinity.tasks import TextClassificationEndpoint, TextClassificationOutput, \
    TextClassificationParams
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import pipeline, AutoTokenizer

class BankingEndpoint(TextClassificationEndpoint):

    __slots__ = ("_pipeline", )

    def __init__(self):
        super().__init__()

        self._pipeline: Optional[ORTModelForSequenceClassification] = None

    def initialize(self, **kwargs):
        print("Initializing")
        model = ORTModelForSequenceClassification.from_pretrained("philschmid/distilbert-onnx-banking77")
        tokenizer = AutoTokenizer.from_pretrained("philschmid/distilbert-onnx-banking77")
        
        self._pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
        print("INITIALIZED")


    def handle(
        self,
        inputs: Union[str, List[str]],
        parameters: TextClassificationParams
    ) -> List[TextClassificationOutput]:
        return self._pipeline(inputs, **parameters)