feat: add TxtSlicesDataset to allow sampling slices from txt file for benchmarking

[hypdeb]

Purpose

Sampling randomly directly from a tokenizer for benchmarking creates data that is not ideal to benchmark when using speculative decoding or expert parallelism.

On the other hand, random datasets are very flexible and offer complete control on the input and output sequence lengths, which is desirable to create reproducible benchmarks.

This PR introduces a new type of benchmarking dataset called TxtSlicesDataset which offers a compromise between the flexibility of a random dataset and the fidelity of a real dataset. It allows sampling slices from a user-provided txt file.

Content

• The implementation of TxtSlicesDataset
• Fixes to typing in datasets.py
• Factored out internal utils from datasets.py in an attempt to bring the file to a more manageable size
• A unit test for the new dataset type

[gemini-code-assist]

Code Review

This pull request introduces TxtSlicesDataset for benchmarking, which samples data from a text file. It also includes significant refactoring by moving utility functions from datasets.py to a new dataset_utils.py file and improving typing throughout. The changes are well-structured. My review focuses on improving the robustness and reproducibility of the new TxtSlicesDataset and its tests. I've pointed out a resource leak in the tests and potential for non-reproducible behavior due to the use of the global random module. I've also identified a missing check that could lead to a crash with certain input files.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hypdeb.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork