label name logic

interfaces/cap_minor.py

@@ -210,11 +210,16 @@ def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
         code = CAP_MIN_NUM_DICT[i]
         prob = probs[i]
 
-        # Minor topic
-        major_code = code // 100
-        major_label = CAP_LABEL_NAMES[major_code]
-        minor_label = CAP_MIN_LABEL_NAMES[code]
-        display = f"[{major_code}] {major_label} [{code}] {minor_label}"
+         if code in CAP_LABEL_NAMES:
+            # Media (major) topic
+            label = CAP_LABEL_NAMES[code]
+            display = f"[{code}] {label}"
+        else:
+            # Minor topic
+            major_code = code // 100
+            major_label = CAP_LABEL_NAMES[major_code]
+            minor_label = CAP_MIN_LABEL_NAMES[code]
+            display = f"[{major_code}] {major_label} [{code}] {minor_label}"
 
         output_pred[display] = prob