File size: 3,045 Bytes
d28af7f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
# 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.
"""
make a general fairseq task for MM pretraining.
"""
import random
from fairseq.tasks import LegacyFairseqTask, register_task
from .task import Task
from .retritask import RetriTask
from ..datasets import FairseqMMDataset
from .. import utils
@register_task("mmtask")
class FairseqMMTask(LegacyFairseqTask):
@staticmethod
def add_args(parser):
# Add some command-line arguments for specifying where the data is
# located and the maximum supported input length.
parser.add_argument(
"taskconfig",
metavar="FILE",
help=("taskconfig to load all configurations" "outside fairseq parser."),
)
@classmethod
def setup_task(cls, args, **kwargs):
return FairseqMMTask(args)
def __init__(self, args):
super().__init__(args)
config = utils.load_config(args)
self.mmtask = Task.config_task(config)
self.mmtask.build_dataset()
self.mmtask.build_model()
self.mmtask.build_loss()
def load_dataset(self, split, **kwargs):
split_map = {
"train": self.mmtask.train_data,
"valid": self.mmtask.val_data,
"test": self.mmtask.test_data,
}
if split not in split_map:
raise ValueError("unknown split type.")
if split_map[split] is not None:
self.datasets[split] = FairseqMMDataset(split_map[split])
def get_batch_iterator(
self,
dataset,
max_tokens=None,
max_sentences=None,
max_positions=None,
ignore_invalid_inputs=False,
required_batch_size_multiple=1,
seed=1,
num_shards=1,
shard_id=0,
num_workers=0,
epoch=1,
data_buffer_size=0,
disable_iterator_cache=False,
skip_remainder_batch=False,
grouped_shuffling=False,
update_epoch_batch_itr=False,
):
random.seed(epoch)
if dataset.mmdataset.split == "train" and isinstance(self.mmtask, RetriTask):
if epoch >= self.mmtask.config.retri_epoch:
if not hasattr(self.mmtask, "retri_dataloader"):
self.mmtask.build_dataloader()
self.mmtask.retrive_candidates(epoch)
return super().get_batch_iterator(
dataset,
max_tokens,
max_sentences,
max_positions,
ignore_invalid_inputs,
required_batch_size_multiple,
seed,
num_shards,
shard_id,
num_workers,
epoch,
data_buffer_size,
disable_iterator_cache,
grouped_shuffling,
update_epoch_batch_itr,
)
@property
def source_dictionary(self):
return None
@property
def target_dictionary(self):
return None
|