Nick Sorros
commited on
Commit
·
f0368c2
1
Parent(s):
33a4c63
Remove the extra list
Browse files- handler.py +1 -1
handler.py
CHANGED
@@ -23,5 +23,5 @@ 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 = [{id2label[label_id]: p} for label_id, p in enumerate(preds[0].tolist()) if p > 0.5]
|
27 |
return prediction
|