normalize

interfaces/cap_minor_media.py

@@ -39,21 +39,15 @@ 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: 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)
+def normalize_probs(probs: dict) -> dict:
+    keys, values = zip(*probs.items())
     values = np.array(values)
-    exp_values = np.exp(values)
+    exp_values = np.exp(values - np.max(values))  # For numerical stability
     sum_exp = np.sum(exp_values)
     normalized = exp_values / sum_exp
-
     return dict(zip(keys, map(float, normalized)))
 
+
 def check_huggingface_path(checkpoint_path: str):
     try:
         hf_api = HfApi(token=HF_TOKEN)