from accelerate import Accelerator
from torch.utils.data import DataLoader
import torch
import webdataset as wds

# ---- Step 1: Initialize Accelerator ---- #
accelerator = Accelerator()
device = accelerator.device

# ---- Step 2: Define WebDataset with IterableDataset ---- #
wds_urls = ["/path/to/shard1.tar", "/path/to/shard2.tar"]  # Add your shard paths here

def custom_collate(batch):
    # Create batch as dict of lists
    collated = {}
    for sample in batch:
        for k, v in sample.items():
            if k not in collated:
                collated[k] = []
            collated[k].append(v)
    # Convert tensors to stacked tensors
    for k in collated:
        if torch.is_tensor(collated[k][0]):
            collated[k] = torch.stack(collated[k])
    return collated

wds_dataset = (
    wds.WebDataset(wds_urls, handler=wds.handlers.warn_and_continue)
    .decode("torchrgb")
    .to_tuple("mp4", "txt", "json")
    .map(lambda sample: {
        "video": sample[0],
        "text": sample[1],
        "meta": sample[2]
    })
)

# ---- Step 3: Define DataLoader ---- #
dataloader = DataLoader(
    wds_dataset,
    batch_size=32,
    num_workers=4,
    pin_memory=True,
    collate_fn=custom_collate
)

# ---- Step 4: Define model and optimizer ---- #
model = torch.nn.Linear(1000, 10)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# ---- Step 5: Prepare model, optimizer, and dataloader ---- #
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

# ---- Step 6: Training Loop ---- #
for epoch in range(10):
    for batch in dataloader:
        batch = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}