Update app.py

app.py

@@ -4,6 +4,8 @@ from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassific
 from safetensors.torch import load_file as safe_load
 
 target_to_ind = {'cs': 0, 'econ': 1, 'eess': 2, 'math': 3, 'phys': 4, 'q-bio': 5, 'q-fin': 6, 'stat': 7}
+target_to_label = {'cs': 'Computer Science', 'econ': 'Economics', 'eess': 'Electrical Engineering and Systems Science', 'math': 'Mathematics', 'phys': 'Physics', 
+                  'q-bio': 'Quantitative Biology', 'q-fin': 'Quantitative Finance', 'stat': 'Statistics'}
 ind_to_target = {ind: target for target, ind in target_to_ind.items()}
 
 
@@ -42,14 +44,24 @@ def get_predict(title: str, abstract: str) -> (str, float, dict):
 title = st.text_area("Title ", "", height=100)
 abstract = st.text_area("Abstract ", "", height=150)
 
-if st.button("Классифицировать", key="manual"):
+
+mode = st.radio("Mode: ", ("Best prediction", "Top 95%"))
+
+if st.button("Get prediction", key="manual"):
     if len(title) == 0:
         st.error("Please, provide paper's title")
     else:
         with st.spinner("Be patient, I'm doing my best"):
             predict = get_predict(title, abstract)
-            
-        st.success(f"Предсказанный тэг: **{predict[0][1]}**")
-        
-        
-model, tokenizer = load_model_and_tokenizer()
+
+        tags = []
+        threshold = 0 if status == "Best prediction" else 0.95
+        sum_p = 0
+        for p, tag in predict:
+            sum_p += p
+            tags.append(target_to_label[tag])
+
+            if sum_p >= threshold:
+                break
+
+        st.succes(f"{'\n'.join(tags)}")