Nick Sorros
commited on
Commit
·
56e1a5f
1
Parent(s):
f0368c2
Add label score in handler
Browse files- handler.py +4 -1
handler.py
CHANGED
@@ -23,5 +23,8 @@ class EndpointHandler():
|
|
23 |
|
24 |
id2label = self.model.config.id2label
|
25 |
|
26 |
-
prediction = [
|
|
|
|
|
|
|
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
|