prettify output for flat pred

interfaces/cap_minor_media.py

@@ -165,16 +165,29 @@ def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
     top_indices = np.argsort(probs)[::-1][:10]
     
     CAP_MIN_MEDIA_LABEL_NAMES = CAP_MEDIA_LABEL_NAMES | CAP_MIN_LABEL_NAMES 
-    output_pred = {
-        f"[{CAP_MIN_MEDIA_NUM_DICT[i]}] {CAP_MIN_MEDIA_LABEL_NAMES[CAP_MIN_MEDIA_NUM_DICT[i]]}": probs[i]
-        for i in top_indices
-    }
-    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>'
+    
+    output_pred = {}
+    for i in top_indices:
+        code = CAP_MIN_MEDIA_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_LABEL_NAMES[code]
+            display = f"[{major_code}] {major_label} [{code}] {minor_label}"
+        
+        output_pred[display] = prob
     
     interpretation_info = """
     ## How to Interpret These Values (Flat Classification)
 
-    This method returns predictions made by a single model. Both media codes and minor topics may appear in the output list. Only the top 10 most confident labels are displayed.
+    This method returns predictions made by a single model. Both media codes and minor topics may appear in the output list. **Only the top 10 most confident labels are displayed**.
     """
     
     return interpretation_info, output_pred, output_info