INNER CODE UNIT · Python

QAEvalDataset

AMontgomerie/question_generator · training/dataset.py:55

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]
        label = random.choice([0, 1])
        if label == 0:
            question, answer = random.choice(self.transforms)(question, answer)
        encoded_data = self.hf_tokenizer(
            text=question,

View source record →