wehub-resource-sync
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import random
|
|
from typing import List
|
|
|
|
from fairseq.data import BaseWrapperDataset, data_utils
|
|
|
|
|
|
class RandomInputDataset(BaseWrapperDataset):
|
|
def __init__(
|
|
self,
|
|
dataset,
|
|
random_input_dataset,
|
|
input_key_path: List[str],
|
|
add_to_input,
|
|
pad_idx,
|
|
):
|
|
super().__init__(dataset)
|
|
self.random_input_dataset = random_input_dataset
|
|
if isinstance(input_key_path, str):
|
|
input_key_path = [input_key_path]
|
|
assert len(input_key_path) > 0
|
|
self.input_key_path = input_key_path
|
|
self.add_to_input = add_to_input
|
|
self.pad_idx = pad_idx
|
|
|
|
def get_target(self, item):
|
|
target_loc = item
|
|
for p in self.input_key_path[:-1]:
|
|
target_loc = target_loc[p]
|
|
return self.input_key_path[-1], target_loc
|
|
|
|
def get_target_value(self, item):
|
|
k, target_loc = self.get_target(item)
|
|
return target_loc[k]
|
|
|
|
def __getitem__(self, index):
|
|
item = self.dataset[index]
|
|
k, target_loc = self.get_target(item)
|
|
target_loc[k] = random.choice(self.random_input_dataset)
|
|
return item
|
|
|
|
def collater(self, samples):
|
|
collated = self.dataset.collater(samples)
|
|
if len(collated) == 0:
|
|
return collated
|
|
indices = set(collated["id"].tolist())
|
|
|
|
random_inputs = data_utils.collate_tokens(
|
|
[self.get_target_value(s) for s in samples if s["id"] in indices],
|
|
pad_idx=self.pad_idx,
|
|
left_pad=False,
|
|
)
|
|
k, target_loc = self.get_target(
|
|
collated if not self.add_to_input else collated["net_input"]
|
|
)
|
|
target_loc[k] = random_inputs
|
|
|
|
return collated
|