removed top n logic

interfaces/cap_minor_media.py

@@ -25,7 +25,7 @@ domains = {
     "media": "media"
 }
 
-
+NUM_TOP_CLASSES = 5
 CAP_MEDIA_CODES = list(CAP_MEDIA_NUM_DICT.values())
 CAP_MIN_CODES = list(CAP_MIN_NUM_DICT.values())
 
@@ -39,11 +39,12 @@ 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):
+def normalize_probs(probs: dict, n=None):
     probs_ = list(probs.values())
-    if len(probs_) > n:
-        probs_.sort(reverse=True)
-        probs_ = probs_[:5]
+    if n:
+        if len(probs_) > n:
+            probs_.sort(reverse=True)
+            probs_ = probs_[:5]
     values = np.array(probs_)
     exp_values = np.exp(values)
     sum_exp = np.sum(exp_values)
@@ -99,7 +100,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, n=5)
+    filtered_probs = normalize_probs(filtered_probs)
         
     output_pred = {
         f"[{major_index_to_id[k]}] {CAP_MEDIA_LABEL_NAMES[major_index_to_id[k]]}": v
@@ -117,7 +118,7 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
         # Restrict to valid minor codes
         valid_indices = [minor_id_to_index[mid] for mid in valid_minor_ids if mid in minor_id_to_index]
         filtered_probs = {minor_index_to_id[i]: float(minor_probs[0][i]) for i in valid_indices}
-        filtered_probs = normalize_probs(filtered_probs, n=5)
+        filtered_probs = normalize_probs(filtered_probs)
         
         output_pred = {
             f"[{top_major_id}] {CAP_MEDIA_LABEL_NAMES[top_major_id]} [{k}] {CAP_MIN_LABEL_NAMES[k]}": v
@@ -161,4 +162,4 @@ demo = gr.Interface(
             gr.Textbox(lines=6, label="Input"),
             gr.Dropdown(languages, label="Language"),
             gr.Dropdown(domains.keys(), label="Domain")],
-    outputs=[gr.Label(num_top_classes=5, label="Output"), gr.Markdown()])
+    outputs=[gr.Label(label="Output"), gr.Markdown()])