| 77 | |
| 78 | |
| 79 | class SummarizationDataset(Dataset): |
| 80 | def __init__( |
| 81 | self, |
| 82 | tokenizer, |
| 83 | data_dir, |
| 84 | type_path="train", |
| 85 | max_source_length=1024, |
| 86 | max_target_length=56, |
| 87 | n_obs=None, |
| 88 | overwrite_cache=False, |
| 89 | prefix="", |
| 90 | ): |
| 91 | super().__init__() |
| 92 | tok_name = tokenizer.__class__.__name__.lower().rstrip("tokenizer") |
| 93 | self.source = encode_file( |
| 94 | tokenizer, |
| 95 | os.path.join(data_dir, type_path + ".source"), |
| 96 | max_source_length, |
| 97 | overwrite_cache=overwrite_cache, |
| 98 | prefix=prefix, |
| 99 | tok_name=tok_name, |
| 100 | ) |
| 101 | tgt_path = os.path.join(data_dir, type_path + ".target") |
| 102 | if hasattr(tokenizer, "set_lang"): |
| 103 | tokenizer.set_lang("ro_RO") # HACK: only applies to mbart |
| 104 | self.target = encode_file( |
| 105 | tokenizer, tgt_path, max_target_length, overwrite_cache=overwrite_cache, tok_name=tok_name |
| 106 | ) |
| 107 | if n_obs is not None: |
| 108 | self.source = self.source[:n_obs] |
| 109 | self.target = self.target[:n_obs] |
| 110 | self.pad_token_id = tokenizer.pad_token_id |
| 111 | |
| 112 | def __len__(self): |
| 113 | return len(self.source) |
| 114 | |
| 115 | def __getitem__(self, index): |
| 116 | source_ids = self.source[index]["input_ids"].squeeze() |
| 117 | target_ids = self.target[index]["input_ids"].squeeze() |
| 118 | src_mask = self.source[index]["attention_mask"].squeeze() |
| 119 | return {"input_ids": source_ids, "attention_mask": src_mask, "decoder_input_ids": target_ids} |
| 120 | |
| 121 | @staticmethod |
| 122 | def trim_seq2seq_batch(batch, pad_token_id): |
| 123 | y = trim_batch(batch["decoder_input_ids"], pad_token_id) |
| 124 | source_ids, source_mask = trim_batch(batch["input_ids"], pad_token_id, attention_mask=batch["attention_mask"]) |
| 125 | return source_ids, source_mask, y |
| 126 | |
| 127 | def collate_fn(self, batch) -> dict: |
| 128 | input_ids = torch.stack([x["input_ids"] for x in batch]) |
| 129 | masks = torch.stack([x["attention_mask"] for x in batch]) |
| 130 | target_ids = torch.stack([x["decoder_input_ids"] for x in batch]) |
| 131 | pad_token_id = self.pad_token_id |
| 132 | y = trim_batch(target_ids, pad_token_id) |
| 133 | source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks) |
| 134 | batch = {"input_ids": source_ids, "attention_mask": source_mask, "decoder_input_ids": y} |
| 135 | return batch |
| 136 |
no outgoing calls