Nick Sorros commited on
Commit
56e1a5f
·
1 Parent(s): f0368c2

Add label score in handler

Browse files
Files changed (1) hide show
  1. handler.py +4 -1
handler.py CHANGED
@@ -23,5 +23,8 @@ class EndpointHandler():
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
 
23
 
24
  id2label = self.model.config.id2label
25
 
26
+ prediction = [
27
+ {"label": id2label[label_id], "score": p}
28
+ for label_id, p in enumerate(preds[0].tolist()) if p > 0.5
29
+ ]
30
  return prediction