Collate missing features

[RicardoDominguez]

Currently the BatchSamplerDataCollatorForSeq2Seq typically requires input_ids, labels, attention_mask, position_ids. For certain uses cases (e.g., datasets of type completion), the latter 3 are straightforwardly derived from input_ids. Therefore, saving to disk datasets with all 4 features can be redundant, and generally requires a lot more disk space (up to 5x).

The changes here implement a MissingFeaturesCollator together with the collate_missing_features config option. If labels, attention_mask, or position_ids are missing from the dataset, it replaces them with the following defaults:

• labels: copy input_ids
• position_ids: [0, 1, 2, ..., len(input_ids)-1]
• attention_mask: [1, 1, 1, ...] with same length as input_ids

I've observed that eagerly filling these missing features does not lead to longer run-times.

[RicardoDominguez]

For instance, pile-cc with input_ids, labels, attention_mask, position_ids, and length is ~1.2TB, and ~250GB with only input_ids and length.

[winglian]

Is there another part that goes with this to optionally have the tokenization step be a bit more sparse for this feature?

[winglian]

Suggested change

	# - attention_mask [1, 1, 1, ...] with same length as input_ids
	# - attention_mask # [1, 1, 1, ...] with same length as input_ids

[winglian]

Looks good so far. I think a unit test on build_collator could be helpful to validate the collator that is returned based on various options, and then some sort of unit test to validate the functionality of MissingFeaturesCollator would be ideal too.

[RicardoDominguez]

I can write something that makes the tokenization step a bit more sparse. How about when all datasets are of type completion?