Nick Sorros
commited on
Commit
·
33a4c63
1
Parent(s):
da5a2e2
Remove need for label binarizer
Browse files- handler.py +2 -4
- requirements.txt +0 -4
handler.py
CHANGED
@@ -1,6 +1,5 @@
|
|
1 |
from typing import Dict, List, Any
|
2 |
from transformers import AutoModel, AutoTokenizer
|
3 |
-
import pickle
|
4 |
|
5 |
|
6 |
class EndpointHandler():
|
@@ -22,8 +21,7 @@ class EndpointHandler():
|
|
22 |
inputs = self.tokenizer(text, padding="max_length")
|
23 |
preds = self.model(input_ids=[inputs["input_ids"]])
|
24 |
|
25 |
-
|
26 |
-
label_binarizer = pickle.loads(f.read())
|
27 |
|
28 |
-
prediction = [[{
|
29 |
return prediction
|
|
|
1 |
from typing import Dict, List, Any
|
2 |
from transformers import AutoModel, AutoTokenizer
|
|
|
3 |
|
4 |
|
5 |
class EndpointHandler():
|
|
|
21 |
inputs = self.tokenizer(text, padding="max_length")
|
22 |
preds = self.model(input_ids=[inputs["input_ids"]])
|
23 |
|
24 |
+
id2label = self.model.config.id2label
|
|
|
25 |
|
26 |
+
prediction = [[{id2label[label_id]: p} for label_id, p in enumerate(preds[0].tolist()) if p > 0.5]]
|
27 |
return prediction
|
requirements.txt
DELETED
@@ -1,4 +0,0 @@
|
|
1 |
-
optimum[onnxruntime]==1.2.3
|
2 |
-
scikit-learn==1.2.2
|
3 |
-
mkl-include
|
4 |
-
mkl
|
|
|
|
|
|
|
|
|
|