Update app.py

app.py

@@ -18,9 +18,20 @@ def process(text):
     """
     Translate incoming text to tokens and classify it
     """
-    pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
+    pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=3)
     result = pipe(text)[0]
-    return result["label"]
+
+    result = sorted(result, key=lambda x: -x["score"])
+
+    cum_score = 0
+    for i, item in enumerate(result):
+        cum_score += item["score"]
+        if cum_score >= 0.95:
+            break
+
+    result = result[: (i + 1)]
+
+    return result
 
 
 tokenizer, model = prepare_model()
@@ -113,4 +124,13 @@ text = "\n".join([title, abstract])
 ## Output
 
 if len(text.strip()) > 0:
-    st.markdown(f"This paper is likely to be from the category **{process(text)}**.")
+    results = process(text)
+    if len(results) == 0:
+        out_text = ""
+    else:
+        out_text = f"This paper is likely to be from the category **{results[0]['label']}** *(score {results[0]['score']:.2f})*."
+        if len(results) > 1:
+            out_text += "\n(Alternative categories are " + " and ".join(
+                [f"{item['label']} *(score {item['score']:.2f})*"]
+            )
+    st.markdown(out_text)