INNER CODE UNIT · Python

__init__

AMontgomerie/question_generator · training/dataset.py:11

    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]
        input_ids, attention_mask = self._encode_text(item.text)

View source record →