INNER CODE UNIT · Python

QGDataset

AMontgomerie/question_generator · training/dataset.py:10

class QGDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        data: datasets.Dataset,
        max_length: int,
        pad_mask_id: int,
        tokenizer: AutoTokenizer
    ) -> None:
        self.data = pd.DataFrame(data)
        self.max_length = max_length
        self.pad_mask_id = pad_mask_id
        self.tokenizer = tokenizer

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, index: int) -> Mapping[str, torch.Tensor]:
        item = self.data.loc[index]

View source record →