import torch


class DummyDataloader:
    def __init__(self, batch_size, vocab_size, seq_length):
        self.batch_size = batch_size
        self.vocab_size = vocab_size
        self.seq_length = seq_length
        self.step = 0

    def generate(self):
        tokens = torch.randint(
            low=0,
            high=self.vocab_size,
            size=(
                self.batch_size,
                self.seq_length,
            ),
        )
        types = torch.randint(
            low=0,
            high=3,
            size=(
                self.batch_size,
                self.seq_length,
            ),
        )
        sentence_order = torch.randint(low=0, high=2, size=(self.batch_size,))
        loss_mask = torch.randint(
            low=0,
            high=2,
            size=(
                self.batch_size,
                self.seq_length,
            ),
        )
        lm_labels = torch.randint(low=0, high=self.vocab_size, size=(self.batch_size, self.seq_length))
        padding_mask = torch.randint(low=0, high=2, size=(self.batch_size, self.seq_length))
        return dict(
            text=tokens,
            types=types,
            is_random=sentence_order,
            loss_mask=loss_mask,
            labels=lm_labels,
            padding_mask=padding_mask,
        )

    def __iter__(self):
        return self

    def __next__(self):
        return self.generate()