fix to correctly specify positives and negatives when passing num_ite…

…rations
huggingface · Jan 13, 2025 · e7f9cf0 · e7f9cf0
1 parent fe30eaa
commit e7f9cf0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/setfit/sampler.py b/src/setfit/sampler.py
@@ -116,8 +116,9 @@ def __init__(
         self.total_neg_pairs = sum(a * b for a, b in combinations(label_counts.values(), 2))
 
         if num_iterations is not None and num_iterations > 0:
-            self.len_pos_pairs = num_iterations * len(self.sentences)
-            self.len_neg_pairs = num_iterations * len(self.sentences)
+            iterations = num_iterations * len(self.sentences)
+            self.len_pos_pairs = int(np.min([self.total_pos_pairs, iterations]))
+            self.len_neg_pairs = int(np.min([self.total_neg_pairs, iterations]))
 
         elif sampling_strategy == SamplingStrategy.UNIQUE:
             self.len_pos_pairs = int(np.min([self.total_pos_pairs, self.max_pos_or_neg]))