refactored normalize_probs

interfaces/cap_minor_media.py

@@ -39,16 +39,20 @@ 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, n=None):
-    probs_ = list(probs.values())
-    if n:
-        if len(probs_) > n:
-            probs_.sort(reverse=True)
-            probs_ = probs_[:5]
-    values = np.array(probs_)
+def normalize_probs(probs: dict, n: int = None) -> dict:
+    if n is not None and len(probs) > n:
+        # Sort items by value in descending order and keep top-n
+        top_items = sorted(probs.items(), key=lambda item: item[1], reverse=True)[:n]
+    else:
+        top_items = list(probs.items())
+
+    keys, values = zip(*top_items)
+    values = np.array(values)
     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)}
+    normalized = exp_values / sum_exp
+
+    return dict(zip(keys, map(float, normalized)))
 
 def check_huggingface_path(checkpoint_path: str):
     try: