feat: Enable cycling through all positive documents in biencoder training #907

[yuhezhang-ai]

What does this PR do ?

This PR updates the biencoder training recipe to utilize all available positive documents for a given query. Previously, only the first positive document was used. Now, the training loop cycles through the list of positive documents across epochs using a modulo operation (e.g., Epoch 0 uses doc 0, Epoch 1 uses doc 1, etc.).

Changelog

• Modified retrieval_dataset.py to accept an epoch argument in the transform function.
• Added update_dataset_epoch helper to update the dataset transform.
• Updated train_biencoder.py to call update_dataset_epoch at the start of each epoch.
• Added unit tests to verify epoch-based positive document selection.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Additional Information

• Close Support multiple positive docs for biencoder training #907

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test c4f83a9

[copy-pr-bot]

/ok to test c4f83a9

@akoumpa, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[akoumpa]

/ok to test f57eaa4

[yuhezhang-ai]

Hi @akoumpa , thanks for triggering CICD test on this PR.

I just noticed there is a newer draft PR (#937) working on the same multi-positive cycling behavior for biencoder training, but using a different implementation:

• My PR: passes the epoch into the dataset transform function and updates the dataset transform at the start of each epoch (update_dataset_epoch). This keeps the existing dataset structure but makes the transform epoch-aware.

• feat: [WIP] Train biencoder with multi pos docs #937: introduces a dataset wrapper that stores an internal _current_epoch field and exposes a set_epoch() API, which changes how the dataset itself is structured.

I wanted to mention this in case it helps decide which approach fits Automodel better. I’m happy to update this PR or align with whichever direction the team prefers.

[akoumpa]

Apologies for the delay @yuhezhang-ai, will let @shan-nvidia to guide the decision further here.

[oliverholworthy]

@yuhezhang-ai Thanks for this PR

I have a suggestion to convert the transform into a stateful class to keep track of the epoch state.

The current implementation calls set_transform() every epoch via update_dataset_epoch(), which recreates the transform closure. A cleaner approach could be to use a stateful transform class something like this:

class RetrievalTransform:
    """Transform for retrieval datasets with epoch-based positive cycling."""
    
    def __init__(self, num_neg_docs: int, corpus_dict: dict):
        self.num_neg_docs = num_neg_docs
        self.corpus_dict = corpus_dict
        self.epoch = 0
    
    def __call__(self, examples):
        return _transform_func(
            examples,
            num_neg_docs=self.num_neg_docs,
            corpus_dict=self.corpus_dict,
            epoch=self.epoch,  # Pass as parameter, don't inject into examples
        )
    
    def set_epoch(self, epoch: int):
        self.epoch = epoch

Then in make_retrieval_dataset():

transform = RetrievalTransform(negative_size, corpus_dict)
dataset.set_transform(transform)  # Called once
dataset.set_epoch = transform.set_epoch  # Expose on dataset

with the following properties:

1. set_transform() called once at creation, not every epoch
2. training loop just call dataset.set_epoch(epoch)
3. Pass epoch as explicit parameter to _transform_func rather than mutating the examples dict
4. No need for update_dataset_epoch() helper - this can be removed
5. State is encapsulated in a class, instead of patched attributes (corpus_dict, num_neg_docs)

Training loop becomes:

if hasattr(self.dataloader.dataset, "set_epoch"):
    self.dataloader.dataset.set_epoch(epoch)