File size: 1,083 Bytes
da45fba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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)
|