add major code mapping

interfaces/cap_minor.py

@@ -8,7 +8,7 @@ from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 from huggingface_hub import HfApi
 
-from label_dicts import CAP_MIN_NUM_DICT, CAP_MIN_LABEL_NAMES
+from label_dicts import CAP_MIN_NUM_DICT, CAP_MIN_LABEL_NAMES, CAP_LABEL_NAMES 
 
 from .utils import is_disk_full
 
@@ -32,6 +32,19 @@ domains = {
     "local government agenda": "localgovernment"
 }
 
+def convert_minor_to_major(results: str) -> str:
+    results_as_text = dict()
+    for i in range(results):
+        prob = probs[i]
+        major_code = CAP_MIN_NUM_DICT[i][:-2]
+        label = CAP_LABEL_NAMES[major_code]
+
+        key = f"[{major_code}] {label}"
+        results_as_text[key] = probs[i]
+
+    return results_as_text
+        
+
 def check_huggingface_path(checkpoint_path: str):
     try:
         hf_api = HfApi(token=HF_TOKEN)
@@ -59,9 +72,10 @@ def predict(text, model_id, tokenizer_id):
         logits = model(**inputs).logits
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
-    output_pred = {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_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])
     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, output_info
+    return output_pred_minor, output_pred_major output_info
 
 def predict_cap(text, language, domain):
     domain = domains[domain]
@@ -80,4 +94,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"), gr.Markdown()])
+    outputs=[gr.Label(num_top_classes=5, label="Output minor"), gr.Label(num_top_classes=5, label="Output major" gr.Markdown()])