text: import functools
  
  
  from torch.utils.data import DataLoader
  from torch.utils.data.dataset import ConcatDataset
  from torch.utils.data.distributed import DistributedSampler
  
  
  from .datamodule_base import BaseDataModule
  from .base_dataset import BaseDataset
  
  
  
  class MTDataModule():
      def __init__(self, _config, dist=False):
          self.dm_dicts = {'coco': CocoCaptionKarpathyDataModule(_config)}
          self.dms = [v for k, v in self.dm_dicts.items()]
  
          self.batch_size = self.dms[0].batch_size
          self.vocab_size = self.dms[0].vocab_size
          self.num_workers = self.dms[0].num_workers
  
          self.dist = dist
  
      def prepare_data(self):
          pass
  
      def setup(self, stage):
          for dm in self.dms:
              dm.setup(stage)
  
          self.train_dataset = ConcatDataset([dm.train_dataset for dm in self.dms])
          self.val_dataset = ConcatDataset([dm.val_dataset for dm in self.dms])
          self.test_dataset = ConcatDataset([dm.test_dataset for dm in self.dms])
          self.tokenizer = self.dms[0].tokenizer
  
          self.collate = functools.partial(
              self.dms[0].train_dataset.collate, mlm_collator=self.dms[0].mlm_collator,
          )
  
          if self.dist:
              self.train_sampler = DistributedSampler(self.train_dataset, shuffle=True)
              self.val_sampler = DistributedSampler(self.val_dataset, shuffle=True)
              self.test_sampler = DistributedSampler(self.test_dataset, shuffle=False)
          else:
              self.train_sampler = None
              self.val_sampler = None
              self.test_sampler = None
  
      def train_dataloader(self):
          loader = DataLoader(
              self.train_dataset,
              batch_size=self.batch_size,
              sampler=self.train_sampler,
              num_workers=self.num_workers,
              collate_fn=self.collate,
          )
          return loader
  
      def val_dataloader(self, batch_size=None):
          loader = DataLoader(
              self.val_dataset,
              batch_size=batch_size if batch_size is not None else self.batch_size,
              sampler=self.val_sampler,
              num_workers=self.num_workers,
              collate_fn=self.collate,
          )
          return loader
  
      def test_dataloader(self):
          loader = DataLoader(
              self.test_dataset,
              batch_size=self.batch_size,
              sampler=self.test_sampler,
              num_workers=self.num_workers,
              collate_fn=self.collate,
          )
          return loader
  
  class CocoCaptionKarpathyDataModule(BaseDataModule):
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
  
      @property
      def dataset_cls(self):
          return CocoCaptionKarpathyDataset
  
      @property
      def dataset_cls_no_false(self):
          return CocoCaptionKarpathyDataset
  
      @property
      def dataset_name(self):
          return "coco"
  
  
  
  class CocoCaptionKarpathyDataset(BaseDataset):
      def __init__(self, *args, split="", **kwargs):
          assert split in ["train", "val", "test"]
          self.split = split
  
          if split == "train":
              names = ["coco_caption_karpathy_train", "coco_caption_karpathy_restval"]
          elif split == "val":
              # names = ["coco_caption_karpathy_val"]
              names = ["coco_caption_karpathy_test"]
          elif split == "test":
              names = ["coco_caption_karpathy_test"]
  
          super().__init__(*args, **kwargs, names=names, text_column_name="caption")
  
      def __getitem__(self, index):
          suite = self.get_suite(index)
  
          if "test" in self.split:
              _index, _question_index = self.index_mapper[index]
              iid = self.table["image_id"][_index].as_py()
              iid = int(iid.split(".")[0].split("_")[-1])
              suite.update({"iid": iid})
  
          return suite