added guide model feature to args and trainer

src/setfit/trainer.py

@@ -1,11 +1,12 @@
+import inspect
 import warnings
 from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple, Union
 
 import evaluate
 import torch
 from datasets import Dataset, DatasetDict
 from packaging.version import parse as parse_version
-from sentence_transformers import SentenceTransformerTrainer, losses
+from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
 from sentence_transformers.losses.BatchHardTripletLoss import BatchHardTripletLossDistanceFunction
 from sentence_transformers.model_card import ModelCardCallback as STModelCardCallback
 from sentence_transformers.training_args import BatchSamplers
@@ -609,7 +610,12 @@ def get_dataset(
                 max_pairs=max_pairs,
             )
             dataset = Dataset.from_generator(data_sampler.__iter__)
-            loss = args.loss(self.model.model_body)
+
+            extra_params = {}
+            if "guide" in inspect.signature(args.loss).parameters:
+                extra_params.update({"guide": SentenceTransformer(args.guide)})
+
+            loss = args.loss(self.model.model_body, **extra_params)
 
         return dataset, loss
 

src/setfit/training_args.py

@@ -161,6 +161,10 @@ class TrainingArguments:
             the case it is "steps", `save_steps` must be a round multiple of `eval_steps`.
 
             </Tip>
+        guide('str', *optional*, defaults to `all-MiniLM-L6-v2`):
+            guide: SentenceTransformer model to guide the in-batch sample selection. Used only with losses that 
+            require guide model such as GISTEmbedLoss.
+
     """
 
     output_dir: str = "checkpoints"
@@ -222,6 +226,8 @@ class TrainingArguments:
     metric_for_best_model: Optional[str] = field(default="embedding_loss", repr=False)
     greater_is_better: bool = field(default=False, repr=False)
 
+    guide: str = "all-MiniLM-L6-v2"
+
     def __post_init__(self) -> None:
         # Set `self.embedding_batch_size` and `self.classifier_batch_size` using values from `self.batch_size`
         if isinstance(self.batch_size, int):