[feat] Improve GroupByLabelBatchSampler

[fpgmaas]

This PR does two things:

Fix the GroupByLabelBatchSampler

I believe this part of code to be incorrect;

    for column_name in valid_label_columns or []:
        if column_name in dataset.column_names:
            labels = dataset["label"]
            break
    else:
        raise ValueError(f"None of the valid_label_columns {valid_label_columns} are in the dataset.")

It seems to try and always use the "label" column, where I expect that the purpose of this code is to find the first value of valid_label_columns that is a valid column in the dataset. This PR fixes that, and make some other small adjustments, for example removing the lines of code mentioned here: #2782

Add unit tests

Current test coverage for sampler.py is 31%, this bumps that up a little bit :) In addition, I think having this unit test can also make it more clear for future developers to understand what the class is aiming to do.

[tomaarsen]

Thanks for this! I do think I'll need changes regarding the self.groups as discussed in that comment itself, but otherwise this is heading in the right direction!

[tomaarsen]

Small nitpick regarding the error

[tomaarsen]

Thanks a bunch for taking care of my review comments. After some nitpicky commit, I think this is all good to go (at least, once the tests go green)!

• Tom Aarsen