LLM continously concat sequence of tokens

quick-dirty-custom-dataloader.py

def custom_dataloader(dataset: Dataset, batch_size=16):
    random_indices = torch.randperm(len(dataset['tokens']) - context_size)

    for idx in range(0, len(random_indices), batch_size):
        x = torch.stack([
            dataset['tokens'][i: i+context_size]
            for i in random_indices[idx: idx+batch_size]
        ])
        y = torch.stack([
            dataset['tokens'][i+1: i+context_size+1]
            for i in random_indices[idx: idx+batch_size]
        ])
        yield x, y