Nick Sorros commited on
Commit
33a4c63
·
1 Parent(s): da5a2e2

Remove need for label binarizer

Browse files
Files changed (2) hide show
  1. handler.py +2 -4
  2. 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
- with open("label_binarizer.pkl", "rb") as f:
26
- label_binarizer = pickle.loads(f.read())
27
 
28
- prediction = [[{label: p} for label, p in zip(label_binarizer.classes_, preds[0].tolist()) if p > 0.5]]
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