INNER CODE UNIT · Python

_mask_label_padding

AMontgomerie/question_generator · training/dataset.py:50

    def _mask_label_padding(self, labels: torch.Tensor) -> torch.Tensor:
        labels[labels == self.tokenizer.pad_token_id] = self.pad_mask_id
        return labels


class QAEvalDataset(torch.utils.data.Dataset):
    def __init__(self, data: datasets.Dataset, max_length: int, tokenizer: AutoTokenizer) -> None:
        self.data = pd.DataFrame(data)
        self.max_length = max_length
        self.transforms = [self.shuffle, self.corrupt]
        self.hf_tokenizer = tokenizer
        self.spacy_tokenizer = en_core_web_sm.load()

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

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

View source record →