Added Infinity endpoint

#1
by mfuntowicz HF staff - opened
Files changed (1) hide show
  1. endpoint.py +31 -0
endpoint.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Optional, Union
2
+
3
+ from infinity.tasks import TextClassificationEndpoint, TextClassificationOutput, \
4
+ TextClassificationParams
5
+ from optimum.onnxruntime import ORTModelForSequenceClassification
6
+ from transformers import pipeline, AutoTokenizer
7
+
8
+ class BankingEndpoint(TextClassificationEndpoint):
9
+
10
+ __slots__ = ("_pipeline", )
11
+
12
+ def __init__(self):
13
+ super().__init__()
14
+
15
+ self._pipeline: Optional[ORTModelForSequenceClassification] = None
16
+
17
+ def initialize(self, **kwargs):
18
+ print("Initializing")
19
+ model = ORTModelForSequenceClassification.from_pretrained("philschmid/distilbert-onnx-banking77")
20
+ tokenizer = AutoTokenizer.from_pretrained("philschmid/distilbert-onnx-banking77")
21
+
22
+ self._pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
23
+ print("INITIALIZED")
24
+
25
+
26
+ def handle(
27
+ self,
28
+ inputs: Union[str, List[str]],
29
+ parameters: TextClassificationParams
30
+ ) -> List[TextClassificationOutput]:
31
+ return self._pipeline(inputs, **parameters)