Implementation of Two Distinct Datasets with HuggingFace Trainer Module

Only if you want to customize updates per gradient , use different optimizers ,conditionally skip batches.

You’d do something like this after you implement the extended class

trainer = MultiDatasetTrainer(
model=model,
args=training_args,
tokenizer=tokenizer,
dataset_a=dataset1,
dataset_b=dataset2,
bs_a=32,
bs_b=128,
train_dataset=None, # Must be None to use overridden method
)
trainer.train()

1 Like