normalize probs (top5)

interfaces/cap_minor_media.py

@@ -39,8 +39,11 @@ for code in CAP_MIN_CODES:
     major_to_minor_map[major_id].append(code)
 major_to_minor_map = dict(major_to_minor_map)
 
-def normalize_probs(probs: dict):
-    values = np.array(list(probs.values()))
+def normalize_probs(probs: dict, n: int):
+    probs = list(probs.values())
+    if len(probs) > n:
+        probs = probs.sort(reverse=True)[:5]
+    values = np.array(probs)
     exp_values = np.exp(values)
     sum_exp = np.sum(exp_values)
     return {k: float(v) for k, v in zip(probs.keys(), exp_values / sum_exp)}
@@ -95,7 +98,7 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
         i: float(major_probs_np[i])
         for i in np.argsort(major_probs_np)[::-1]
     }
-    filtered_probs = normalize_probs(filtered_probs)
+    filtered_probs = normalize_probs(filtered_probs, n=5)
         
     output_pred = {
         f"[{major_index_to_id[k]}] {CAP_MEDIA_LABEL_NAMES[major_index_to_id[k]]}": v