chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,62 @@
+# 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