update inference

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import arxiv
-from inference import classify, load_model  # Replace with your actual module
+from inference import top_95_labels, load_model  # Replace with your actual module
 
 @st.cache_resource
 def get_model():
@@ -41,6 +41,8 @@ elif input_method == "Manual Input":
 
 # Classification and output
 if title and abstract:
-    category = classify(pipe, class2name, title, abstract)
+    categories = top_95_labels(title, abstract)  # Assuming classify returns a list of labels
     st.write(f"### Title: {title}")
-    st.write(f"**Predicted Category:** {category}")
+    st.write("**Predicted Categories:**")
+    for category in categories:
+        st.write(f"- {category}")

inference.py

@@ -16,9 +16,16 @@ def load_model(path2chkpt: str, path2mapping: str):
     return pipe, class2name
 
 
-def classify(pipe: Pipeline, class2name: dict[str, str], title: str, abstract: str):
+def top_95_labels(pipe: Pipeline, class2name: dict[str, str], title: str, abstract: str):
     inputs = ".".join([title, abstract])
-    class_code = pipe(inputs)[0]["label"]
+    result = pipe(inputs, top_k=20)
     
-    return class2name[class_code]
+    proba = 0
+    labels = []
+    i = 0
+    while proba < 0.95:
+        proba += result[i]["score"]
+        labels.append(result[i]["label"])
+        i += 1
+    return [class2name[label] for label in labels]