| |
| |
| |
| |
|
|
| import torch |
|
|
| from fairseq.data import BaseWrapperDataset, data_utils |
|
|
|
|
| class AddClassTargetDataset(BaseWrapperDataset): |
| def __init__( |
| self, |
| dataset, |
| labels, |
| multi_class, |
| num_classes=None, |
| label_indices=None, |
| add_to_input=True, |
| ): |
| super().__init__(dataset) |
|
|
| self.label_indices = label_indices |
| self.labels = labels |
| self.multi_class = multi_class |
| self.add_to_input = add_to_input |
| if num_classes is None and multi_class: |
| assert self.label_indices is not None |
| num_classes = len(self.label_indices) |
|
|
| self.num_classes = num_classes |
|
|
| def __getitem__(self, index): |
| item = self.dataset[index] |
| item_labels = self.labels[index] |
| if self.multi_class: |
| item["label"] = torch.zeros(self.num_classes) |
| for il in item_labels: |
| if self.label_indices is not None: |
| il = self.label_indices[il] |
| item["label"][il] = 1.0 |
| else: |
| item["label"] = torch.tensor( |
| self.labels[index] |
| if self.label_indices is None |
| else self.label_indices[self.labels[index]] |
| ) |
|
|
| return item |
|
|
| 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] |
| collated["label"] = torch.stack(target, dim=0) |
|
|
| if self.add_to_input: |
| collated["net_input"]["label"] = collated["label"] |
|
|
| return collated |
|
|
