Enforce passing item_loader when customizing underlying storage format

[tchaton]

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Until now, we had a hack where 1D tensor would be handled differentely and stored as contiguous array. I have seen several users complaining about this magic and unexpected behaviour.

WARNING: This PR is a breaking change for LLM using the TokensLoader. Now, we would need to pass the item_loader to the optimize or Cache directly to inform the underlying storage needs to be handled differently.

If no item_loader is passed during the optimization, this default to the Pytree handler.

For LLM tokens, here is the breaking API change.

Before

from litdata.streaming.item_loader import TokensLoader

optimize(...)

dataset = StreamingDataset(item_loader=TokensLoader(...))

Now

from litdata.streaming.item_loader import TokensLoader

optimize(..., item_loader=TokensLoader())

dataset = StreamingDataset(item_loader=TokensLoader(...))

Fixes #294

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[lantiga]

Much cleaner and more robust. IMO the change is not confusing, one just has to use item loader on both sides.
What happens when I (erroneously) omit it?

[tchaton]

What happens when I (erroneously) omit it?

It would break and tell you the wrong item loader has been provided. Ok, waiting on @awaelchli review and ideas ;)

[AugustDev]

TokensLoader accepts block_size. What if each of my sample is a dictionary where some of the fields are token sequences. Should block_size be length of the longest sequence?

[tchaton]

Hey @AugustDev. TokensLoader doesn't work on dictionary. It needs a 1D tensor.