|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
