Made sure all labels are represented in the training set

embeddings-benchmark · May 8, 2024 · 96312c7 · 96312c7
1 parent 49afeb1
commit 96312c7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mteb/abstasks/AbsTaskMultilabelClassification.py b/mteb/abstasks/AbsTaskMultilabelClassification.py
@@ -168,8 +168,11 @@ def _undersample_data_indices(self, y, samples_per_label, idxs=None):
             idxs = np.arange(len(y))
         np.random.shuffle(idxs)
         label_counter = defaultdict(int)
+        unique_labels = set(itertools.chain.from_iterable(y))
         for i in idxs:
-            if any((label_counter[label] < samples_per_label) for label in y[i]):
+            if any(
+                (label_counter[label] < samples_per_label) for label in unique_labels
+            ):
                 sample_indices.append(i)
                 for label in y[i]:
                     label_counter[label] += 1