Update app.py

app.py

@@ -13,8 +13,20 @@ all_label_names = list(model.config.id2label.values())
 
 def predict(text):
     probs = expit(model(**tokenizer([text], return_tensors="pt", padding=True)).logits.detach().numpy())
-    d = {i: float(np.round(j, 2)) for i, j in zip(all_label_names, probs[0])}
-    print(d)
+    # can't use numpy for whatever reason
+    probs = [float(np.round(i, 2)) for i in probs[0]]
+    # break out issue, harm, sentiment, feeling
+    zipped_list = list(zip(all_label_names, probs))
+    issues = [(i, j) for i, j in zipped_list if i.startswith('issue')]
+    feelings = [(i, j) for i, j in zipped_list if i.startswith('feeling')]
+    harm = [(i, j) for i, j in zipped_list if i.startswith('harm')]
+    # keep tops for each one
+    issues = sorted(issues)[::-1][:3]
+    feelings = sorted(feelings)[::-1][:3]
+    harm = sorted(harm)[::-1][:1]
+    # top is the combo of these
+    top = issues + feelings + harm
+    d = {i: j for i, j in harm}
     return d
 
 iface = gr.Interface(