Add training option unique_pairs

scripts/setfit/run_fewshot.py

@@ -55,6 +55,7 @@ def parse_args():
     parser.add_argument("--override_results", default=False, action="store_true")
     parser.add_argument("--keep_body_frozen", default=False, action="store_true")
     parser.add_argument("--add_data_augmentation", default=False)
+    parser.add_argument("--unique_pairs", type=bool, default=False)
 
     args = parser.parse_args()
 
@@ -143,6 +144,7 @@ def main():
                 batch_size=args.batch_size,
                 num_epochs=args.num_epochs,
                 num_iterations=args.num_iterations,
+                unique_pairs=args.unique_pairs,
             )
             if args.classifier == "pytorch":
                 trainer.freeze()

src/setfit/modeling.py

@@ -1,5 +1,6 @@
 import os
 from dataclasses import dataclass
+from itertools import combinations, combinations_with_replacement, zip_longest
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
@@ -633,30 +634,87 @@ def forward(self, sentence_features, labels=None, mask=None):
         return loss
 
 
-def sentence_pairs_generation(sentences, labels, pairs):
-    # Initialize two empty lists to hold the (sentence, sentence) pairs and
-    # labels to indicate if a pair is positive or negative
+def positive_sentence_pairs_generate(
+    sentences: "np.ndarray[str]", labels: "np.ndarray[int]", max_pairs: int, unique_pairs: bool
+) -> List[InputExample]:
+    """Generates all unique or upto a max no. of combinations of positive sentence pairs.
 
-    num_classes = np.unique(labels)
-    idx = [np.where(labels == i)[0] for i in num_classes]
+    Samples positive combinations of sentences (without replacement) and maximises
+        sampling of different classes in the pairs being generated.
 
-    for first_idx in range(len(sentences)):
-        current_sentence = sentences[first_idx]
-        label = labels[first_idx]
-        second_idx = np.random.choice(idx[np.where(num_classes == label)[0][0]])
-        positive_sentence = sentences[second_idx]
-        # Prepare a positive pair and update the sentences and labels
-        # lists, respectively
-        pairs.append(InputExample(texts=[current_sentence, positive_sentence], label=1.0))
-
-        negative_idx = np.where(labels != label)[0]
-        negative_sentence = sentences[np.random.choice(negative_idx)]
-        # Prepare a negative pair of sentences and update our lists
-        pairs.append(InputExample(texts=[current_sentence, negative_sentence], label=0.0))
-    # Return a 2-tuple of our sentence pairs and labels
+    Args:
+        sentences: an array of sentences
+        labels: an array of the label_id for each sentence in `sentences`
+        max_pairs: returns when this many pairs are generated
+        unique_pairs: if true will return sentences if all unique combinations,
+            before max_pairs count is reached
+
+    Returns:
+        List of positive sentence pairs (upto the no. of unique_pairs or max_pairs)
+    """
+    pairs = []
+    while True:
+        positive_combinators = []
+        for _label in np.unique(labels):
+            label_sentences = sentences[np.where(labels == _label)]
+            positive_combinators.append(combinations_with_replacement(label_sentences, 2))
+
+        for pos_pairs in zip_longest(*positive_combinators):
+            for pos_pair in pos_pairs:
+                if pos_pair is not None:
+                    pairs.append(InputExample(texts=[*pos_pair], label=1.0))
+                    if len(pairs) == max_pairs:
+                        return pairs
+
+        if unique_pairs:
+            break
+    logger.warning(f"** All ({len(pairs):,}) positive unique pairs generated")
+    return pairs
+
+
+def negative_sentence_pairs_generate(
+    sentences: "np.ndarray[str]", labels: "np.ndarray[int]", max_pairs: int, unique_pairs: bool
+) -> List[InputExample]:
+    """Generates all or upto a max sample no. of negative combinations.
+
+    Randomly samples negative combinations of sentences (without replacement)
+
+    Args:
+        sentences: an array of sentences
+        labels: an array of the label_id for each sentence in `sentences`
+        max_pairs: returns when this many pairs are generated
+        unique_pairs: if true will return sentences if all unique combinations,
+            before max_pairs count is reached
+
+    Returns:
+        List of negative sentence pairs (upto the no. of unique_pairs or max_pairs)
+    """
+    pairs = []
+    while True:
+        sent_labels = [(sent, label) for sent, label in zip(sentences, labels)]
+        for (_sentence, _label), (sentence, label) in combinations(sent_labels, 2):
+            if _label != label:
+                pairs.append(InputExample(texts=[_sentence, sentence], label=0.0))
+                if len(pairs) == max_pairs:
+                    return pairs
+        if unique_pairs:
+            break
+    logger.warning(f"** All ({len(pairs):,}) negative unique pairs generated")
     return pairs
 
 
+def sentence_pairs_generation(
+    sentences: "np.ndarray[str]", labels: "np.ndarray[int]", max_pairs: int, unique_pairs: bool
+) -> List[InputExample]:
+    max_pos_pairs = max_pairs / 2
+    positive_pairs = positive_sentence_pairs_generate(sentences, labels, max_pos_pairs, unique_pairs)
+
+    max_neg_pairs = max_pairs - len(positive_pairs)
+    negative_pairs = negative_sentence_pairs_generate(sentences, labels, max_neg_pairs, unique_pairs)
+
+    return positive_pairs + negative_pairs
+
+
 def sentence_pairs_generation_multilabel(sentences, labels, pairs):
     # Initialize two empty lists to hold the (sentence, sentence) pairs and
     # labels to indicate if a pair is positive or negative

src/setfit/trainer.py

@@ -94,6 +94,7 @@ def __init__(
         distance_metric: Callable = BatchHardTripletLossDistanceFunction.cosine_distance,
         margin: float = 0.25,
         samples_per_label: int = 2,
+        unique_pairs: bool = False,
     ):
         if (warmup_proportion < 0.0) or (warmup_proportion > 1.0):
             raise ValueError(
@@ -116,6 +117,7 @@ def __init__(
         self.distance_metric = distance_metric
         self.margin = margin
         self.samples_per_label = samples_per_label
+        self.unique_pairs = unique_pairs
 
         if model is None:
             if model_init is not None:
@@ -349,17 +351,17 @@ def train(
 
                 train_steps = len(train_dataloader) * self.num_epochs
             else:
-                train_examples = []
-
-                for _ in range(self.num_iterations):
-                    if self.model.multi_target_strategy is not None:
+                if self.model.multi_target_strategy is not None:  # TODO
+                    train_examples = []
+                    for _ in range(self.num_iterations):
                         train_examples = sentence_pairs_generation_multilabel(
                             np.array(x_train), np.array(y_train), train_examples
                         )
-                    else:
-                        train_examples = sentence_pairs_generation(
-                            np.array(x_train), np.array(y_train), train_examples
-                        )
+                else:
+                    max_pairs = self.num_iterations * len(x_train) * 2
+                    train_examples = sentence_pairs_generation(
+                        np.array(x_train), np.array(y_train), max_pairs, self.unique_pairs
+                    )
 
                 train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=batch_size)
                 train_loss = self.loss_class(self.model.model_body)