Support for the differentiable head

scripts/setfit/run_fewshot.py

@@ -51,6 +51,7 @@ def parse_args():
     parser.add_argument("--is_dev_set", type=bool, default=False)
     parser.add_argument("--is_test_set", type=bool, default=False)
     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)
 
     args = parser.parse_args()
@@ -105,7 +106,14 @@ def main():
                 continue
 
             # Load model
-            model = SetFitModel.from_pretrained(args.model)
+            if args.classifier == "pytorch":
+                model = SetFitModel.from_pretrained(
+                    args.model,
+                    use_differentiable_head=True,
+                    head_params={"out_features": len(set(train_data["label"]))},
+                )
+            else:
+                model = SetFitModel.from_pretrained(args.model)
             model.model_body.max_seq_length = args.max_seq_length
             if args.add_normalization_layer:
                 model.model_body._modules["2"] = models.Normalize()
@@ -121,7 +129,19 @@ def main():
                 num_epochs=args.num_epochs,
                 num_iterations=args.num_iterations,
             )
-            trainer.train()
+            if args.classifier == "pytorch":
+                trainer.freeze()
+                trainer.train()
+                trainer.unfreeze(keep_body_frozen=args.keep_body_frozen)
+                trainer.train(
+                    num_epochs=25,
+                    body_learning_rate=1e-5,
+                    learning_rate=args.lr,  # recommend: 1e-2
+                    l2_weight=0.0,
+                    batch_size=args.batch_size,
+                )
+            else:
+                trainer.train()
 
             # Evaluate the model on the test data
             metrics = trainer.evaluate()

src/setfit/__init__.py

@@ -1,4 +1,4 @@
 __version__ = "0.4.0.dev0"
 
-from .modeling import SetFitModel
+from .modeling import SetFitHead, SetFitModel
 from .trainer import SetFitTrainer

src/setfit/data.py

@@ -1,9 +1,16 @@
-from typing import Dict, List
+from typing import TYPE_CHECKING, Dict, List, Tuple
 
 import pandas as pd
+import torch
 from datasets import Dataset, DatasetDict
+from torch.utils.data import Dataset as TorchDataset
 
 
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+
+
+TokenizerOutput = Dict[str, List[int]]
 SEEDS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
 SAMPLE_SIZES = [2, 4, 8, 16, 32, 64]
 
@@ -189,3 +196,71 @@ def add_templated_examples(
             dataset = dataset.add_item(example)
 
     return dataset
+
+
+class SetFitDataset(TorchDataset):
+    """SetFitDataset
+
+    A dataset for training the differentiable head on text classification.
+
+    Args:
+        x (`List[str]`):
+            A list of input data as texts that will be fed into `SetFitModel`.
+        y (`List[int]`):
+            A list of input data's labels.
+        tokenizer (`PreTrainedTokenizerBase`):
+            The tokenizer from `SetFitModel`'s body.
+        max_length (`int`, defaults to `32`):
+            The maximum token length a tokenizer can generate.
+            Will pad or truncate tokens when the number of tokens for a text is either smaller or larger than this value.
+    """
+
+    def __init__(
+        self,
+        x: List[str],
+        y: List[int],
+        tokenizer: "PreTrainedTokenizerBase",
+        max_length: int = 32,
+    ) -> None:
+        assert len(x) == len(y)
+
+        self.x = x
+        self.y = y
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self) -> int:
+        return len(self.x)
+
+    def __getitem__(self, idx: int) -> Tuple[TokenizerOutput, int]:
+        feature = self.tokenizer(
+            self.x[idx],
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            return_token_type_ids=True,
+        )
+        label = self.y[idx]
+
+        return feature, label
+
+    @staticmethod
+    def collate_fn(batch):
+        features = {
+            "input_ids": [],
+            "attention_mask": [],
+            "token_type_ids": [],
+        }
+        labels = []
+        for feature, label in batch:
+            features["input_ids"].append(feature["input_ids"])
+            features["attention_mask"].append(feature["attention_mask"])
+            features["token_type_ids"].append(feature["token_type_ids"])
+            labels.append(label)
+
+        # convert to tensors
+        features = {k: torch.Tensor(v).int() for k, v in features.items()}
+        labels = torch.Tensor(labels).long()
+
+        return features, labels