| |
| |
| |
| |
|
|
| import torch |
|
|
| from . import BaseWrapperDataset, data_utils |
|
|
|
|
| class AddTargetDataset(BaseWrapperDataset): |
| def __init__( |
| self, |
| dataset, |
| labels, |
| pad, |
| eos, |
| batch_targets, |
| process_label=None, |
| add_to_input=False, |
| ): |
| super().__init__(dataset) |
| self.labels = labels |
| self.batch_targets = batch_targets |
| self.pad = pad |
| self.eos = eos |
| self.process_label = process_label |
| self.add_to_input = add_to_input |
|
|
| def get_label(self, index): |
| return ( |
| self.labels[index] |
| if self.process_label is None |
| else self.process_label(self.labels[index]) |
| ) |
|
|
| def __getitem__(self, index): |
| item = self.dataset[index] |
| item["label"] = self.get_label(index) |
| return item |
|
|
| def size(self, index): |
| sz = self.dataset.size(index) |
| own_sz = len(self.get_label(index)) |
| return (sz, own_sz) |
|
|
| def collater(self, samples): |
| collated = self.dataset.collater(samples) |
| if len(collated) == 0: |
| return collated |
| indices = set(collated["id"].tolist()) |
| target = [s["label"] for s in samples if s["id"] in indices] |
|
|
| if self.batch_targets: |
| collated["target_lengths"] = torch.LongTensor([len(t) for t in target]) |
| target = data_utils.collate_tokens(target, pad_idx=self.pad, left_pad=False) |
| collated["ntokens"] = collated["target_lengths"].sum().item() |
| else: |
| collated["ntokens"] = sum([len(t) for t in target]) |
|
|
| collated["target"] = target |
|
|
| if self.add_to_input: |
| eos = target.new_full((target.size(0), 1), self.eos) |
| collated["target"] = torch.cat([target, eos], dim=-1).long() |
| collated["net_input"]["prev_output_tokens"] = torch.cat( |
| [eos, target], dim=-1 |
| ).long() |
| collated["ntokens"] += target.size(0) |
| return collated |
|
|
