Option for `(Distributed)LengthGroupedSampler` to treat groups as a hard constraint
Author: mbforbes
Pull request: no

# 🚀 Feature request

An option for `(Distributed)LengthGroupedSampler` to treat groups as a hard constraint. I.e., all batches returned will have exactly the same length. (Some straggler batches will then have a smaller batch size.)

## Motivation

I asked a [question](https://discuss.huggingface.co/t/multiple-choice-with-variable-number-of-choices/8607) on the forums about using a classification model to do multiple choice with a variable number of choices.

The simplest implementation I can see using HF Transformers, though it's a bit of a hack, is to use `--group_by_length` and set `--length_column_name` to be the number of choices. That way, the `1` dimension, which tells the model the number of multiple choice options, is consistent throughout a batch.

This _almost_ works. The issue is that `(Distributed)LengthGroupedSampler` is a soft constraint, so some batches still end up with multiple "lengths" (choices).

During training, I check each batch in the collator, and simply throw away samples that don't have the same number of choices.

The issue is that during evaluation, I realized I can't skip over any examples. In comparing methods, we need to report results on the entire evaluation set.

I totally understand if you'd rather not add this feature for the purpose of what is, admittedly, a hack. If that's the case, I'd greatly appreciate any advice on how you'd run a multiple choice model with a variable number of choices!

Comments:
This is definitely a hack and very narrow use case, so it's unlikely we will add this feature ;-).
For your purpose, I think you need to rewrite a new model head, starting from the code of `XxxForMultipleChoice`, where `Xxx` is the model you are using.

Oh interesting, I hadn't considered changing the model, since it's already agnostic to the number of answers. But I suppose it could do the final bit of data preprocessing in the model itself!

I'm a bit surprised that there are no multiple choice datasets with a variable number of answers. Maybe this is something we'll see NLP tackle soon :-) 

After thinking about this a bit more, it seemed like a custom (batch) sampler was the most elegant way to accomplish this.

Since multiple choice models are already agnostic to the number of choices, we just need something to feed it batches where each batch has a consistent number of choices.

Here's an example implementation of a single-process batch sampler that accomplishes this. It simply groups by feature upon construction, then provides iterators that yield batches with a particular value of that feature in common.

```python
class FeatureGroupedBatchSampler(Sampler):
    """Yields a batch of indices at a time, with the hard constraint that all indices
    will have the same value for `feature`.

    From pytorch docs:
    "Mutually exclusive with batch_size, shuffle, sampler, and drop_last."

    NOTE: shuffle, drop_last not yet implemented. Will if needed.
    """

    dataset: datasets.Dataset
    feature: str
    batch_size: int
    val2idxes: Dict[Any, List[int]]
    num_batches: int

    def __init__(
        self, dataset: datasets.Dataset, feature: str, batch_size: int
    ) -> None:
        if not isinstance(dataset, datasets.Dataset):
            raise ValueError("`dataset` must be a (HuggingFace) datasets.Dataset")
        if feature not in dataset.features:
            raise ValueError(f"Feature '{feature}' must exist on dataset")

        self.dataset = dataset
        self.feature = feature
        self.batch_size = batch_size

        val2idxes = defaultdict(list)
        for i, val in enumerate(dataset[self.feature]):
            val2idxes[val].append(i)
        self.val2idxes = val2idxes
        # NOTE: Only the indices (dict's values) are ever used. Could remove features
        # (dict's keys) entirely if desired and store as List[List[int]].

        # Cache the number of batches so we don't need to recompute it on calls to
        # __len__().
        num_batches = 0
        for idxes in self.val2idxes.values():
            for start in range(0, len(idxes), self.batch_size):
                num_batches += 1
        self.num_batches = num_batches

    def __iter__(self) -> Iterator[List[int]]:
        """Yields a batch of indicies at a time. Maximum of `self.batch_size`, but
        always with identical values for `self.feature`.
        """
        for idxes in self.val2idxes.values():
            for start in range(0, len(idxes), self.batch_size):
                yield idxes[start : start + self.batch_size]

    def __len__(self) -> int:
        """At least the way HF Transformers uses this, it means number of *batches*,
        not number of *instances* (since this is sent as a batch_sampler)."""
        # return len(self.dataset)  # this would be for num. instances
        return self.num_batches  # num. batches
```

I use the above during evaluation by overriding the`Trainer`'s `get_eval_dataloader()` and loading it, providing the "number of choices" column of the dataset as the feature to group by.

However, there's one issue: it doesn't seem like 🤗 Transformers supports a custom batch sampler. It expects the batch size to be known by the data loader, which is not the case if a custom batch sampler is provided.

https://github.com/huggingface/transformers/blob/1fec32adc6a4840123d5ec5ff5cf419c02342b5a/src/transformers/trainer.py#L2172

This leads to a crash a few lines later:

https://github.com/huggingface/transformers/blob/1fec32adc6a4840123d5ec5ff5cf419c02342b5a/src/transformers/trainer.py#L2217

I have a one-line fix that I'll propose in a PR, which is to simply use the observed batch size, which was calculated a few lines before:

https://github.com/huggingface/transformers/blob/1fec32adc6a4840123d5ec5ff5cf419c02342b5a/src/transformers/trainer.py#L2208


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.