change design

interfaces/cap_minor.py

@@ -32,24 +32,15 @@ domains = {
     "local government agenda": "localgovernment"
 }
 
-def convert_minor_to_major(results, probs):
-    results_as_text = dict()
-    for i in results:
-        prob = probs[i]
-        major_code = str(CAP_MIN_NUM_DICT[i])[:-2]
-
-        if major_code == "99":
-            major_code = "999"
-        
-        label = CAP_LABEL_NAMES[int(major_code)]
+def convert_minor_to_major(minor_topic):
+    major_code = str(minor_topic)[:-2]
 
-        key = f"[{major_code}] {label}"
-        if key in results_as_text:
-            results_as_text[key] += probs[i]
-        else:
-            results_as_text[key] = probs[i]
+    if major_code == "99":
+        major_code = "999"
+    
+    label = CAP_LABEL_NAMES[int(major_code)]
 
-    return results_as_text
+    return label
         
 
 def check_huggingface_path(checkpoint_path: str):
@@ -79,10 +70,9 @@ def predict(text, model_id, tokenizer_id):
         logits = model(**inputs).logits
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
-    output_pred_minor = {f"[{CAP_MIN_NUM_DICT[i]}] {CAP_MIN_LABEL_NAMES[CAP_MIN_NUM_DICT[i]]}": probs[i] for i in np.argsort(probs)[::-1]}
-    output_pred_major = convert_minor_to_major(np.argsort(probs)[::-1], probs)
+    output_pred = {f"[{CAP_MIN_NUM_DICT[i]}] {convert_minor_to_major(CAP_MIN_NUM_DICT[i])} - {CAP_MIN_LABEL_NAMES[CAP_MIN_NUM_DICT[i]]}": probs[i] for i in np.argsort(probs)[::-1]}
     output_info = f'<p style="text-align: center; display: block">Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.</p>'
-    return output_pred_minor, output_pred_major, output_info
+    return output_pred, output_info
 
 def predict_cap(text, language, domain):
     domain = domains[domain]
@@ -101,4 +91,4 @@ demo = gr.Interface(
     inputs=[gr.Textbox(lines=6, label="Input"),
             gr.Dropdown(languages, label="Language"),
             gr.Dropdown(domains.keys(), label="Domain")],
-    outputs=[gr.Label(num_top_classes=5, label="Output minor"), gr.Label(num_top_classes=5, label="Output major"), gr.Markdown()])
+    outputs=[gr.Label(num_top_classes=5, label="Output"), gr.Markdown()])