Add n_documents option for HF datasets

[danbraunai-apollo]

Add n_documents option for HF datasets

Description

• Add n_documents as an attribute of HFDatasetConfig. This governs how many documents to load from the dataset, before splitting into n_samples lots of n_ctx-length samples.
• Add seed attribute to HFDatasetConfig. This governs which random selection of n_ctx-length samples are selected from the larger set of documents

How Has This Been Tested?

• Add test_data.test_invalid_hf_dataset_config for checking the pydantic validation of HFDatasetConfig
• Add TestTokenizeDataset which checks various properties of rib.loader.tokenize_dataset, that now also takes in n_samples and seed arguments.

Does this PR introduce a breaking change?

Yes. return_set_n_samples has changed names to n_samples. Also, previously when return_set_n_samples was used in a pythia config, it actually specified the number of documents loaded from the dataset. It now specifies the number of n_ctx-length samples.

Note that:

• tinystories has ~235 toks / document.
• pile-10k, which we use for pythia, has ~1555 toks / document

[nix-apollo]

Some other dataset related improvements that could be added (feel free to skip)

• [x] Support for the tiny stories test set, which is titled "validation" instead of "test"
• [ ] Caching downloaded data (and models?) in some folder on the ssd drive. Best case we can set transformers to offline mode and shave a second off of imports (nice for fast tests!)

[danbraunai-apollo]

• [ ] Caching downloaded data (and models?) in some folder on the ssd drive. Best case we can set transformers to offline mode and shave a second off of imports (nice for fast tests!)

Yeah this would be nice. I haven't looked at "offline mode" before. Not going to look into it now, but made an issue for it.

[danbraunai-apollo]

Several changes since last review. Best for you to have a look. Most notably, as discussed, the test_stochastic_basis_tinystories should be made more robust to different seeds and/or n_samples in the dataset config.