Allow passing a custom batch sampler to the trainer #3162
Conversation
| @@ -77,21 +97,17 @@ def __init__( | |||
| dataset: Dataset, | |||
| batch_size: int, | |||
| drop_last: bool, | |||
| valid_label_columns: list[str] = None, | |||
| generator: torch.Generator = None, | |||
| seed: int = 0, | |||
There was a problem hiding this comment.
I believe there was an issue with the typing and usage of the seed parameter, given that 0 is falsy, it would be ignored if it was passed as a seed.
| if self._batch_sampler: | ||
| return self._batch_sampler |
There was a problem hiding this comment.
I like the changes in the sampler.py, but I'm not sure if this is the best option.
This here requires the user to initialize the Batch Sampler themselves, which prevents Sentence Transformers from updating the dataset(s). This breaks the prompts feature here:
sentence-transformers/sentence_transformers/trainer.py
Lines 284 to 291 in c68bf68
Additionally, it prevents multi-dataset training setups, because there's only one batch sampler possible.
So, I think a good solution is for the batch_sampler argument to be 1) a class (not instance) that subclasses DefaultBatchSampler or 2) a function that returns a subclass of DefaultBatchSampler given dataset, batch_size, drop_last, valid_label_columns, generator, seed, and *args & **kwargs.
What do you think?
resolves #3152
Also, I believe there was an issue with the typing and usage of the seed parameter, given that 0 is falsy, it would be ignored if it was passed as a seed.
Missing: