label names, num dicts

interfaces/cap_minor.py

@@ -129,7 +129,7 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
     filtered_probs = normalize_probs(filtered_probs)
 
     output_pred = {
-        f"[{major_index_to_id[k]}] {CAP_MEDIA_LABEL_NAMES[major_index_to_id[k]]}": v
+        f"[{major_index_to_id[k]}] {CAP_LABEL_NAMES[major_index_to_id[k]]}": v
         for k, v in sorted(
             filtered_probs.items(), key=lambda item: item[1], reverse=True
         )
@@ -161,7 +161,7 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
         print(filtered_probs)  # debug
 
         output_pred = {
-            f"[{top_major_id}] {CAP_MEDIA_LABEL_NAMES[top_major_id]} [{k}] {CAP_MIN_LABEL_NAMES[k]}": v
+            f"[{top_major_id}] {CAP_LABEL_NAMES[top_major_id]} [{k}] {CAP_MIN_LABEL_NAMES[k]}": v
             for k, v in sorted(
                 filtered_probs.items(), key=lambda item: item[1], reverse=True
             )
@@ -205,23 +205,16 @@ def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     top_indices = np.argsort(probs)[::-1][:10]
 
-    CAP_MIN_MEDIA_LABEL_NAMES = CAP_MEDIA_LABEL_NAMES | CAP_MIN_LABEL_NAMES
-
     output_pred = {}
     for i in top_indices:
-        code = CAP_MIN_MEDIA_NUM_DICT[i]
+        code = CAP_MIN_NUM_DICT[i]
         prob = probs[i]
 
-        if code in CAP_MEDIA_LABEL_NAMES:
-            # Media (major) topic
-            label = CAP_MEDIA_LABEL_NAMES[code]
-            display = f"[{code}] {label}"
-        else:
-            # Minor topic
-            major_code = code // 100
-            major_label = CAP_MEDIA_LABEL_NAMES[major_code]
-            minor_label = CAP_MIN_LABEL_NAMES[code]
-            display = f"[{major_code}] {major_label} [{code}] {minor_label}"
+        # 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