Update level_classifier_tool_2.py

level_classifier_tool_2.py

@@ -138,8 +138,6 @@ def _aggregate_sims(
     """
     if sims.numel() == 0:
         return float("nan")
-    if agg == "mean":
-        return float(sims.mean().item())
     if agg == "max":
         return float(sims.max().item())
     if agg == "topk_mean":
@@ -238,7 +236,7 @@ def classify_levels_phrases(
                 top_contribs[lvl] = []
                 continue
             sims = (q_emb @ embs.T).squeeze(0)  # cosine sim due to L2 norm
-            scores[lvl] = _aggregate_sims(sims, agg, topk)
+            scores[lvl] = _aggregate_sims(sims, max, topk)
             if return_phrase_matches:
                 k = min(5, sims.numel())
                 vals, idxs = torch.topk(sims, k)