feat(litdata): Add grouping functionality in StreamingRawDataset

[bhimrazy]

What does this PR do ?

Enhance the StreamingRawDataset class with grouping capabilities in the setup method, allowing for flexible item structuring for users. Improve asynchronous handling for batch downloads and update tests to ensure functionality. Remove unused imports and streamline the code.

Fixes #662

Follow up to #652

Usage Example

from litdata.streaming.raw_dataset import StreamingRawDataset, FileMetadata
from torch.utils.data import DataLoader
from typing import Union

class CustomStreamingRawDataset(StreamingRawDataset):
    def setup(self, files: list[FileMetadata]) -> Union[list[FileMetadata], list[list[FileMetadata]]]:
       # Customize grouping logic here.
       # For example: return pairs like [[image_1, mask_1], [image_2, mask_2], ...]
        return files

# Initialize the streaming raw dataset from S3 path
dataset = CustomStreamingRawDataset("s3://bucket/files/")

[codecov]

Codecov Report

❌ Patch coverage is 94.28571% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 84%. Comparing base (f88a139) to head (9e1f8ee).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@         Coverage Diff         @@
##           main   #665   +/-   ##
===================================
  Coverage    83%    84%           
===================================
  Files        50     50           
  Lines      7124   7137   +13     
===================================
+ Hits       5935   5960   +25     
+ Misses     1189   1177   -12     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[tchaton]

Mind adding example in the README ?

[bhimrazy]

Mind adding example in the README ?

Sure! Just updated.
My bad — I missed adding it earlier. Thanks for the reminder!

[tchaton]

Looks great !