hunk.code_window: def load_checkpoint_to_cpu(path, arg_overrides=None):
       """Loads a checkpoint to CPU (with upgrading for backward compatibility)."""
       with PathManager.open(path, "rb") as f:
           state = torch.load(
               f, map_location=lambda s, l: default_restore_location(s, "cpu")
           )
  
hunk.labels: keep keep keep keep replace keep keep keep keep
hunk.after_edit:     with open(PathManager.get_local_path(path), "rb") as f:
  
hunk.file_path: fairseq/checkpoint_utils.py
hunk.type: replace
file: # 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 itertools
  import json
  import logging
  import os
  from collections import OrderedDict, defaultdict
  
  from fairseq import utils
  from fairseq.data import (
      AppendTokenDataset,
      ConcatDataset,
      Dictionary,
      LanguagePairDataset,
      PrependTokenDataset,
      SampledMultiDataset,
      SampledMultiEpochDataset,
      StripTokenDataset,
      TransformEosLangPairDataset,
      TruncateDataset,
      data_utils,
      indexed_dataset,
  )
  from fairseq.data.multilingual.multilingual_utils import (
      EncoderLangtok,
      LangTokSpec,
      LangTokStyle,
      augment_dictionary,
      get_lang_tok,
  )
  from fairseq.data.multilingual.sampled_multi_dataset import CollateFormat
  from fairseq.file_io import PathManager
  from fairseq.utils import FileContentsAction, csv_str_list, eval_str_dict
  
  
  logger = logging.getLogger(__name__)
  
  
  def _lang_id(dic: Dictionary, lang: str):
      """Return language ID index."""
      idx = dic.index(lang)
      assert idx != dic.unk_index, "cannot find language ID for lang {}".format(lang)
      return idx
  
  
  def load_sampling_weights(from_file):
      with open(from_file) as f:
          weights = json.load(f)
      return weights
  
  
  class MultilingualDatasetManager(object):
      def __init__(self, args, lang_pairs, langs, dicts, sampling_method):
          super().__init__()
          self.args = args
          self.seed = args.seed
          self.lang_pairs = lang_pairs
          self.langs = langs
          self.dicts = dicts
          self.lang_dict = self.create_lang_dictionary(self.langs)
          self.sampling_method = sampling_method
          self.sampling_scheduler = None
          self._has_sharded_data = False
          self._num_shards_dict = {}
          self._training_data_sizes = defaultdict(lambda: {})
  
      @classmethod
      def setup_data_manager(cls, args, lang_pairs, langs, dicts, sampling_method):
          return MultilingualDatasetManager(
              args, lang_pairs, langs, dicts, sampling_method
          )
  
      @staticmethod
      def add_args(parser):
          parser.add_argument(
              "data",
              help="colon separated path to data directories list, \
                              will be iterated upon during epochs in round-robin manner",
              action=FileContentsAction,
          )
          parser.add_argument(
              "--langs",
              default=None,
              type=csv_str_list,
              help="a list of languages comma sperated languages which can appear in lang-pairs; "
              "note that the ordering determines language token IDs",
              action=FileContentsAction,
          )
          parser.add_argument(
              "--lang-dict",
              default=None,
              type=str,
              help="an external file which contains a list of "
              "languages which can appear in lang-pairs; "
              "note that the ordering determines language token IDs; "
              "--langs and --lang-dict are two exclusive options",
          )
          parser.add_argument(
              "--lang-tok-style",
              default=LangTokStyle.multilingual.value,
              type=str,
              choices=[LangTokStyle.multilingual.value, LangTokStyle.mbart.value],
              help="language token styles",
          )
  
          parser.add_argument(
              "--load-alignments",
              action="store_true",
              help="load the binarized alignments",
          )
          parser.add_argument(
              "--left-pad-source",
              default="True",
              type=str,
              metavar="BOOL",
              help="pad the source on the left",
          )
          parser.add_argument(
              "--left-pad-target",
              default="False",
              type=str,
              metavar="BOOL",
              help="pad the target on the left",
          )
          parser.add_argument(
              "--max-source-positions",
              default=1024,
              type=int,
              metavar="N",
              help="max number of tokens in the source sequence",
          )
          parser.add_argument(
              "--max-target-positions",
              default=1024,
              type=int,
              metavar="N",
              help="max number of tokens in the target sequence",
          )
          parser.add_argument(
              "--upsample-primary",
              default=1,
              type=int,
              help="amount to upsample primary dataset",
          )
          parser.add_argument(
              "--truncate-source",
              action="store_true",
              default=False,
              help="truncate source to max-source-positions",
          )
          parser.add_argument(
              "--encoder-langtok",
              default=None,
              type=str,
              choices=[EncoderLangtok.src.value, EncoderLangtok.tgt.value],
              metavar="SRCTGT",
              help="prepend to the beginning of source sentence the source or target "
              "language token. (src/tgt)",
          )
          parser.add_argument(
              "--decoder-langtok",
              action="store_true",
              help="prepend to the beginning of target sentence the target language token",
          )
          parser.add_argument(
              "--lang-tok-replacing-bos-eos", action="store_true", default=False
          )
          parser.add_argument(
              "--enable-lang-ids",
              default=False,
              action="store_true",
              help="whether to include language IDs in samples",
          )
          parser.add_argument(
              "--enable-reservsed-directions-shared-datasets",
              default=False,
              action="store_true",
              help="whether to allow datasets be used in reversed directions",
          )
  
          parser.add_argument(
              "--extra-data",
              help='a dictionary of data name to this path, \
                              e.g. {"mined", path_to_mined_data, "denoised": path_to_denoised_data}',
              type=lambda uf: eval_str_dict(uf, type=str),
              default=None,
          )
          parser.add_argument(
              "--extra-lang-pairs",
              help='a dictionary of data name to the language pairs they serve, \
                              e.g. {"mined": comma-separated-lang-pairs, "denoised":  comma-separated-lang-pairs}',
              type=lambda uf: eval_str_dict(uf, type=str),
              default=None,
          )
          parser.add_argument(
              "--langtoks-specs",
              help='a list of comma separated data types that a set of language tokens to be specialized for, \
                              e.g. "main,dae,mined". There will be a set of language tokens added to the vocab to \
                              distinguish languages in different training data types. If not specified, default language \
                              tokens per languages will be added',
              default=LangTokSpec.main.value,
              type=csv_str_list,
          )
          parser.add_argument(
              "--langtoks",
              help='a dictionary of how to add language tokens, \
                              e.g. {"mined": (None, "tgt"), "mono_dae": ("src.dae", "tgt"), "main": \
                              ("src", "tgt")}, or {"mined": ("src.mined", "tgt")}',
              default=None,
              type=lambda uf: eval_str_dict(uf, type=str),
          )
          parser.add_argument(
              "--sampling-weights-from-file",
              help='a file contain a python dictionary of how to sample data sets, \
                                  e.g. { "main:en_XX-es_XX": 0.2, "mined:en_XX-pt_XX": 0.5, \
                                      "mono_dae:es_XX-es_XX: 0.3, "main:en_xx-fr_XX": 0.8 }',
              default=None,
              type=str,
          )
          parser.add_argument(
              "--sampling-weights",
              help='a dictionary of how to sample data sets, \
                              e.g. { "main:en_XX-es_XX": 0.2, "mined:en_XX-pt_XX": 0.5, \
                                     "mono_dae:es_XX-es_XX: 0.3, "main:en_xx-fr_XX": 0.8 }',
              default=None,
              type=lambda uf: eval_str_dict(uf, type=str),
          )
          parser.add_argument(
              "--virtual-epoch-size",
              default=1000000,
              type=int,
              help="virtual epoch size to speed up data loading",
          )
          parser.add_argument(
              "--virtual-data-size",
              default=None,
              type=int,
              help="virtual data size of the whole joint dataset to speed"
              "up data loading and have specific dynamic sampling strategy interval",
          )
  
      @classmethod
      def load_langs(cls, args, **kwargs):
          if args.lang_dict and args.langs:
              raise ValueError("--langs and --lang-dict can not both be specified")
          if args.lang_dict is None and args.langs is None:
              logger.warning(
                  "External language dictionary is not provided; "
                  "use lang-pairs to infer the set of supported languages. "
                  "The language ordering is not stable which might cause "
                  "misalignment in pretraining and finetuning."
              )
              # infer from lang_pairs as it is
              langs = list(
                  {x for lang_pair in args.lang_pairs for x in lang_pair.split("-")}
              )
              langs = sorted(langs)
              logger.info(f"inferred language list: {langs}")
          elif args.lang_dict:
              with PathManager.open(args.lang_dict, "r", encoding="utf-8") as f:
                  langs = [lang.strip() for lang in f.readlines() if lang.strip()]
                  logger.info(
                      f"loaded language list from {args.lang_dict} as they are ordered in file"
                  )
          elif args.langs:
              langs = args.langs
              logger.info(
                  f"parsed the language list as they are ordered in the option: {langs}"
              )
          return langs
  
      def has_sharded_data(self, split):
          return self._has_sharded_data and split == getattr(
              self.args, "train_subset", None
          )
  
      def _shared_collater(self):
          return not (self.args.extra_data and "mono_dae" in self.args.extra_data) and (
              not self.args.lang_tok_replacing_bos_eos
          )
  
      @classmethod
      def prepare(cls, load_dictionary, args, **kargs):
          args.left_pad_source = utils.eval_bool(args.left_pad_source)
          args.left_pad_target = utils.eval_bool(args.left_pad_target)
  
          if not hasattr(args, "shuffle_instance"):
              args.shuffle_instance = False
          if args.langtoks is None:
              args.langtoks = {}
          if "main" not in args.langtoks:
              src_langtok_spec = args.encoder_langtok if args.encoder_langtok else None
              tgt_langtok_spec = "tgt" if args.decoder_langtok else None
              args.langtoks["main"] = (src_langtok_spec, tgt_langtok_spec)
  
          def check_langs(langs, pairs):
              messages = []
              for src, tgt in pairs:
                  if src not in langs or tgt not in langs:
                      messages.append(
                          f"language pair {src}-{tgt} contains languages "
                          "that are not in the language dictionary"
                      )
              if len(messages) > 0:
                  raise ValueError(" ".join(messages) + f"; langs: {langs}")
  
          if args.lang_pairs is None:
              raise ValueError(
                  "--lang-pairs is required. List all the language pairs in the training objective."
              )
          if isinstance(args.lang_pairs, str):
              args.lang_pairs = args.lang_pairs.split(",")
          if args.source_lang is not None or args.target_lang is not None:
              training = False
          else:
              training = True
          language_list = cls.load_langs(args, **kargs)
          check_langs(
              language_list,
              (
                  [p.split("-") for p in args.lang_pairs]
                  if training
                  else [(args.source_lang, args.target_lang)]
              ),
          )
  
          # load dictionaries
          if training:
              extra_lang_pairs = (
                  list(
                      {p for _, v in args.extra_lang_pairs.items() for p in v.split(",")}
                  )
                  if args.extra_lang_pairs
                  else []
              )
              langs_to_load_dicts = sorted(
                  {x for p in args.lang_pairs + extra_lang_pairs for x in p.split("-")}
              )
          else:
              langs_to_load_dicts = sorted([args.source_lang, args.target_lang])
  
          dicts = OrderedDict()
          paths = utils.split_paths(args.data)
          assert len(paths) > 0
          for lang in langs_to_load_dicts:
              dicts[lang] = load_dictionary(
                  os.path.join(paths[0], "dict.{}.txt".format(lang))
              )
              augment_dictionary(
                  dictionary=dicts[lang],
                  language_list=language_list,
                  lang_tok_style=args.lang_tok_style,
                  langtoks_specs=args.langtoks_specs,
                  extra_data=args.extra_data,
              )
              if len(dicts) > 0:
                  assert dicts[lang].pad() == dicts[langs_to_load_dicts[0]].pad()
                  assert dicts[lang].eos() == dicts[langs_to_load_dicts[0]].eos()
                  assert dicts[lang].unk() == dicts[langs_to_load_dicts[0]].unk()
              logger.info("[{}] dictionary: {} types".format(lang, len(dicts[lang])))
          return language_list, dicts, training
  
      @classmethod
      def create_lang_dictionary(cls, langs):
          unk = "<unk>"
          # hack to remove symbols other than unk as they are not needed by lang dict
          lang_dict = Dictionary(pad=unk, eos=unk, unk=unk, bos=unk)
          for lang in langs:
              lang_dict.add_symbol(lang)
          return lang_dict
  
      @classmethod
      def get_langtok_index(cls, lang_tok, dic):
          idx = dic.index(lang_tok)
          assert (
              idx != dic.unk_index
          ), "cannot find language token {} in the dictionary".format(lang_tok)
          return idx
  
      def get_encoder_langtok(self, src_lang, tgt_lang, spec=None):
          if spec is None:
              return None
          if spec and spec.startswith("src"):
              if src_lang is None:
                  return None
              langtok = get_lang_tok(
                  lang=src_lang, lang_tok_style=self.args.lang_tok_style, spec=spec
              )
          else:
              if tgt_lang is None:
                  return None
              langtok = get_lang_tok(
                  lang=tgt_lang, lang_tok_style=self.args.lang_tok_style, spec=spec
              )
          return self.get_langtok_index(
              langtok, self.dicts[src_lang if src_lang else tgt_lang]
          )
  
      def get_decoder_langtok(self, tgt_lang, spec=None):
          if spec is None:
              return None
          langtok = get_lang_tok(
              lang=tgt_lang, lang_tok_style=self.args.lang_tok_style, spec=spec
          )
          return self.get_langtok_index(langtok, self.dicts[tgt_lang])
  
      @classmethod
      def load_data(cls, path, vdict, impl):
          dataset = data_utils.load_indexed_dataset(path, vdict, impl)
          return dataset
  
      @classmethod
      def split_exists(cls, split, src, tgt, lang, data_path, dataset_impl):
          filename = os.path.join(data_path, "{}.{}-{}.{}".format(split, src, tgt, lang))
          return indexed_dataset.dataset_exists(filename, impl=dataset_impl)
  
      def load_lang_dataset(
          self,
          data_path,
          split,
          src,
          src_dict,
          tgt,
          tgt_dict,
          combine,
          dataset_impl,
          upsample_primary,
          max_source_positions,
          prepend_bos=False,
          load_alignments=False,
          truncate_source=False,
      ):
  
          src_datasets = []
          tgt_datasets = []
  
          for k in itertools.count():
              split_k = split + (str(k) if k > 0 else "")
  
              # infer langcode
              if self.split_exists(split_k, src, tgt, src, data_path, dataset_impl):
                  prefix = os.path.join(data_path, "{}.{}-{}.".format(split_k, src, tgt))
              elif self.split_exists(split_k, tgt, src, src, data_path, dataset_impl):
                  prefix = os.path.join(data_path, "{}.{}-{}.".format(split_k, tgt, src))
              else:
                  if k > 0:
                      break
                  else:
                      logger.error(
                          f"Dataset not found: {data_path}, {split_k}, {src}, {tgt}"
                      )
                      raise FileNotFoundError(
                          "Dataset not found: {} ({})".format(split, data_path)
                      )
  
              src_dataset = self.load_data(prefix + src, src_dict, dataset_impl)
              if truncate_source:
                  src_dataset = AppendTokenDataset(
                      TruncateDataset(
                          StripTokenDataset(src_dataset, src_dict.eos()),
                          max_source_positions - 1,
                      ),
                      src_dict.eos(),
                  )
              src_datasets.append(src_dataset)
              tgt_datasets.append(self.load_data(prefix + tgt, tgt_dict, dataset_impl))
  
              logger.info(
                  "{} {} {}-{} {} examples".format(
                      data_path, split_k, src, tgt, len(src_datasets[-1])
                  )
              )
  
              if not combine:
                  break
  
          assert len(src_datasets) == len(tgt_datasets)
  
          if len(src_datasets) == 1:
              src_dataset, tgt_dataset = src_datasets[0], tgt_datasets[0]
          else:
              sample_ratios = [1] * len(src_datasets)
              sample_ratios[0] = upsample_primary
              src_dataset = ConcatDataset(src_datasets, sample_ratios)
              tgt_dataset = ConcatDataset(tgt_datasets, sample_ratios)
  
          if prepend_bos:
              assert hasattr(src_dict, "bos_index") and hasattr(tgt_dict, "bos_index")
              src_dataset = PrependTokenDataset(src_dataset, src_dict.bos())
              tgt_dataset = PrependTokenDataset(tgt_dataset, tgt_dict.bos())
  
          align_dataset = None
          if load_alignments:
              align_path = os.path.join(
                  data_path, "{}.align.{}-{}".format(split, src, tgt)
              )
              if indexed_dataset.dataset_exists(align_path, impl=dataset_impl):
                  align_dataset = data_utils.load_indexed_dataset(
                      align_path, None, dataset_impl
                  )
  
          return src_dataset, tgt_dataset, align_dataset
  
      def load_langpair_dataset(
          self,
          data_path,
          split,
          src,
          src_dict,
          tgt,
          tgt_dict,
          combine,
          dataset_impl,
          upsample_primary,
          left_pad_source,
          left_pad_target,
          max_source_positions,
          max_target_positions,
          prepend_bos=False,
          load_alignments=False,
          truncate_source=False,
          src_dataset_transform_func=lambda dataset: dataset,
          tgt_dataset_transform_func=lambda dataset: dataset,
          src_lang_id=None,
          tgt_lang_id=None,
          langpairs_sharing_datasets=None,
      ):
          norm_direction = "-".join(sorted([src, tgt]))
          if langpairs_sharing_datasets is not None:
              src_dataset = langpairs_sharing_datasets.get(
                  (data_path, split, norm_direction, src), "NotInCache"
              )
              tgt_dataset = langpairs_sharing_datasets.get(
                  (data_path, split, norm_direction, tgt), "NotInCache"
              )
              align_dataset = langpairs_sharing_datasets.get(
                  (data_path, split, norm_direction, src, tgt), "NotInCache"
              )
  
          # a hack: any one is not in cache, we need to reload them
          if (
              langpairs_sharing_datasets is None
              or src_dataset == "NotInCache"
              or tgt_dataset == "NotInCache"
              or align_dataset == "NotInCache"
              or split != getattr(self.args, "train_subset", None)
          ):
              # source and target datasets can be reused in reversed directions to save memory
              # reversed directions of valid and test data will not share source and target datasets
              src_dataset, tgt_dataset, align_dataset = self.load_lang_dataset(
                  data_path,
                  split,
                  src,
                  src_dict,
                  tgt,
                  tgt_dict,
                  combine,
                  dataset_impl,
                  upsample_primary,
                  max_source_positions=max_source_positions,
                  prepend_bos=prepend_bos,
                  load_alignments=load_alignments,
                  truncate_source=truncate_source,
              )
              src_dataset = src_dataset_transform_func(src_dataset)
              tgt_dataset = tgt_dataset_transform_func(tgt_dataset)
              if langpairs_sharing_datasets is not None:
                  langpairs_sharing_datasets[
                      (data_path, split, norm_direction, src)
                  ] = src_dataset
                  langpairs_sharing_datasets[
                      (data_path, split, norm_direction, tgt)
                  ] = tgt_dataset
                  langpairs_sharing_datasets[
                      (data_path, split, norm_direction, src, tgt)
                  ] = align_dataset
                  if align_dataset is None:
                      # no align data so flag the reverse direction as well in sharing
                      langpairs_sharing_datasets[
                          (data_path, split, norm_direction, tgt, src)
                      ] = align_dataset
          else:
              logger.info(
                  f"Reusing source and target datasets of [{split}] {tgt}-{src} for reversed direction: "
                  f"[{split}] {src}-{tgt}: src length={len(src_dataset)}; tgt length={len(tgt_dataset)}"
              )
  
          return LanguagePairDataset(
              src_dataset,
              src_dataset.sizes,
              src_dict,
              tgt_dataset,
              tgt_dataset.sizes if tgt_dataset is not None else None,
              tgt_dict,
              left_pad_source=left_pad_source,
              left_pad_target=left_pad_target,
              align_dataset=align_dataset,
              src_lang_id=src_lang_id,
              tgt_lang_id=tgt_lang_id,
          )
  
      def src_dataset_tranform_func(self, src_lang, tgt_lang, dataset, spec=None):
          if self.args.lang_tok_replacing_bos_eos:
              # it is handled by self.alter_dataset_langtok
              # TODO: Unifiy with alter_dataset_langtok
              return dataset
          if spec is None:
              return dataset
          tok = self.get_encoder_langtok(src_lang, tgt_lang, spec)
          if tok:
              return PrependTokenDataset(dataset, tok)
          return dataset
  
      def tgt_dataset_tranform_func(self, source_lang, target_lang, dataset, spec=None):
          if dataset is None:
              # note that target dataset can be None during inference time
              return None
          if self.args.lang_tok_replacing_bos_eos:
              # TODO: Unifiy with alter_dataset_langtok
              # It is handled by self.alter_dataset_langtok.
              # The complication in self.alter_dataset_langtok
              # makes a unified framework difficult.
              return dataset
          # if not self.args.decoder_langtok:
          if not spec:
              return dataset
          tok = self.get_decoder_langtok(target_lang, spec)
          if tok:
              return PrependTokenDataset(dataset, tok)
          return dataset
  
      def alter_dataset_langtok(
          self,
          lang_pair_dataset,
          src_eos=None,
          src_lang=None,
          tgt_eos=None,
          tgt_lang=None,
          src_langtok_spec=None,
          tgt_langtok_spec=None,
      ):
          if src_langtok_spec is None and tgt_langtok_spec is None:
              return lang_pair_dataset
  
          new_src_eos = None
          if (
              src_langtok_spec is not None
              and src_eos is not None
              and (src_lang is not None or tgt_lang is not None)
          ):
              new_src_eos = self.get_encoder_langtok(src_lang, tgt_lang, src_langtok_spec)
          else:
              src_eos = None
  
          new_tgt_bos = None
          if tgt_langtok_spec and tgt_eos is not None and tgt_lang is not None:
              new_tgt_bos = self.get_decoder_langtok(tgt_lang, tgt_langtok_spec)
          else:
              tgt_eos = None
  
          return TransformEosLangPairDataset(
              lang_pair_dataset,
              src_eos=src_eos,
              new_src_eos=new_src_eos,
              tgt_bos=tgt_eos,
              new_tgt_bos=new_tgt_bos,
          )
  
      def load_a_dataset(
          self,
          split,
          data_path,
          src,
          src_dict,
          tgt,
          tgt_dict,
          combine,
          prepend_bos=False,
          langpairs_sharing_datasets=None,
          data_category=None,
          **extra_kwargs,
      ):
          dataset_impl = self.args.dataset_impl
          upsample_primary = self.args.upsample_primary
          left_pad_source = self.args.left_pad_source
          left_pad_target = self.args.left_pad_target
          max_source_positions = self.args.max_source_positions
          max_target_positions = self.args.max_target_positions
          load_alignments = self.args.load_alignments
          truncate_source = self.args.truncate_source
          src_dataset_transform_func = self.src_dataset_tranform_func
          tgt_dataset_transform_func = self.tgt_dataset_tranform_func
          enable_lang_ids = self.args.enable_lang_ids
          lang_dictionary = self.lang_dict
          src_langtok_spec, tgt_langtok_spec = extra_kwargs["langtok_spec"]
  
          src_langtok = self.get_encoder_langtok(src, tgt, src_langtok_spec)
          tgt_langtok = self.get_decoder_langtok(tgt, tgt_langtok_spec)
          logger.info(
              f"{data_category}:{src}-{tgt} src_langtok: {src_langtok}; tgt_langtok: {tgt_langtok}"
          )
  
          langpair_ds = self.load_langpair_dataset(
              data_path,
              split,
              src,
              src_dict,
              tgt,
              tgt_dict,
              combine,
              dataset_impl,
              upsample_primary,
              left_pad_source,
              left_pad_target,
              max_source_positions,
              max_target_positions,
              prepend_bos,
              load_alignments,
              truncate_source,
              src_dataset_transform_func=lambda dataset: src_dataset_transform_func(
                  src, tgt, dataset, src_langtok_spec
              ),
              tgt_dataset_transform_func=lambda dataset: tgt_dataset_transform_func(
                  src, tgt, dataset, tgt_langtok_spec
              ),
              src_lang_id=_lang_id(lang_dictionary, src)
              if enable_lang_ids and lang_dictionary is not None
              else None,
              tgt_lang_id=_lang_id(lang_dictionary, tgt)
              if enable_lang_ids and lang_dictionary is not None
              else None,
              langpairs_sharing_datasets=langpairs_sharing_datasets,
          )
          # TODO: handle modified lang toks for mined data and dae data
          if self.args.lang_tok_replacing_bos_eos:
              ds = self.alter_dataset_langtok(
                  langpair_ds,
                  src_eos=self.dicts[src if src else tgt].eos(),
                  src_lang=src,
                  tgt_eos=self.dicts[tgt].eos(),
                  tgt_lang=tgt,
                  src_langtok_spec=src_langtok_spec,
                  tgt_langtok_spec=tgt_langtok_spec,
              )
          else:
              ds = langpair_ds
          return ds
  
      def load_split_langpair_datasets(self, split, data_param_list):
          datasets = []
          langpairs_sharing_datasets = (
              {} if self.args.enable_reservsed_directions_shared_datasets else None
          )
          for param in data_param_list:
              ds = self.load_a_dataset(
                  split=split,
                  langpairs_sharing_datasets=langpairs_sharing_datasets,
                  **param,
              )
              datasets.append(ds)
          return datasets
  
      def get_data_paths_and_lang_pairs(self, split):
          datapaths = {"main": self.args.data}
          lang_pairs = {"main": self.lang_pairs}
          if split == getattr(self.args, "train_subset", None):
              # only training data can have extra data and extra language pairs
              if self.args.extra_data:
                  extra_datapaths = self.args.extra_data
                  datapaths.update(extra_datapaths)
              if self.args.extra_lang_pairs:
                  extra_lang_pairs = {
                      k: v.split(",") for k, v in self.args.extra_lang_pairs.items()
                  }
                  lang_pairs.update(extra_lang_pairs)
          return datapaths, lang_pairs
  
      @classmethod
      def get_dataset_key(cls, data_category, src, tgt):
          return f"{data_category}:{src}-{tgt}"
  
      @classmethod
      def _get_shard_num_dict(cls, split, paths):
          shards = defaultdict(int)
          for path in paths:
              files = PathManager.ls(path)
              for f in files:
                  if f.startswith(split) and f.endswith(".idx"):
                      # idx files of the form "{split}.{src}-{tgt}.{lang}.idx"
                      direction = f.split(".")[-3]
                      shards[direction] += 1
          # each direction has two '.idx' files
          # one for source language and one for target language, so:
          return {k: v // 2 for k, v in shards.items()}
  
      def get_split_num_data_shards(self, split):
          if split in self._num_shards_dict:
              return self._num_shards_dict[split]
          num_shards_dict = {}
          data_paths, lang_pairs = self.get_data_paths_and_lang_pairs(split)
  
          for data_category, paths in data_paths.items():
              if data_category not in lang_pairs:
                  continue
              paths = utils.split_paths(paths)
              shards_dict = self._get_shard_num_dict(split, paths)
              lang_dirs = [
                  lang_pair.split("-") for lang_pair in lang_pairs[data_category]
              ]
              lang_dirs = [x if len(x) > 1 else (x[0], x[0]) for x in lang_dirs]
              for src, tgt in lang_dirs:
                  key = self.get_dataset_key(data_category, src, tgt)
                  if "mono_" in data_category:
                      # monolingual data requires tgt only
                      assert src is None or src == tgt, (
                          f"error: src={src}, "
                          "tgt={tgt} for data_category={data_category}"
                      )
                      num_shards_dict[key] = shards_dict[tgt]
                  else:
                      if f"{src}-{tgt}" in shards_dict:
                          num_shards_dict[key] = shards_dict[f"{src}-{tgt}"]
                      elif f"{tgt}-{src}" in shards_dict:
                          # follow the fairseq tradition to use reversed direction data if it is not available
                          num_shards_dict[key] = shards_dict[f"{tgt}-{src}"]
          self._num_shards_dict[split] = num_shards_dict
          logger.info(f"[{split}] num of shards: {num_shards_dict}")
          return num_shards_dict
  
      @classmethod
      def get_shard_id(cls, num_shards, epoch, shard_epoch=None):
          shard = epoch if shard_epoch is None else shard_epoch
          shard = (shard - 1) % num_shards
          return shard
  
      def get_split_data_path(self, paths, epoch, shard_epoch, num_shards):
          path = paths[self.get_shard_id(num_shards, epoch, shard_epoch)]
          return path
  
      def get_split_data_param_list(self, split, epoch, shard_epoch=None):
          # TODO: to extend with extra datasets and keys and loop over different shard data paths
          param_list = []
          data_paths, lang_pairs = self.get_data_paths_and_lang_pairs(split)
          logger.info(f"langtoks settings: {self.args.langtoks}")
          split_num_shards_dict = self.get_split_num_data_shards(split)
          for data_category, paths in data_paths.items():
              if data_category not in lang_pairs:
                  continue
              paths = utils.split_paths(paths)
              assert len(paths) > 0
              if len(paths) > 1:
                  self._has_sharded_data = True
              if split != getattr(self.args, "train_subset", None):
                  # if not training data set, use the first shard for valid and test
                  paths = paths[:1]
  
              if data_category in self.args.langtoks:
                  lang_tok_spec = self.args.langtoks[data_category]
              else:
                  # default to None
                  lang_tok_spec = (None, None)
  
              # infer langcode
              lang_dirs = [
                  lang_pair.split("-") for lang_pair in lang_pairs[data_category]
              ]
              lang_dirs = [x if len(x) > 1 else (x[0], x[0]) for x in lang_dirs]
              for src, tgt in lang_dirs:
                  assert src is not None or data_category == "mono_dae", (
                      f"error: src={src}, " "tgt={tgt} for data_category={data_category}"
                  )
                  # logger.info(f"preparing param for {data_category}: {src} - {tgt}")
                  key = self.get_dataset_key(data_category, src, tgt)
                  data_path = self.get_split_data_path(
                      paths, epoch, shard_epoch, split_num_shards_dict[key]
                  )
                  param_list.append(
                      {
                          "key": key,
                          "data_path": data_path,
                          "split": split,
                          "src": src,
                          "src_dict": self.dicts[src]
                          if src and data_category != "mono_dae"
                          else None,
                          "tgt": tgt,
                          "tgt_dict": self.dicts[tgt],
                          "data_category": data_category,
                          "langtok_spec": lang_tok_spec,
                      }
                  )
          return param_list
  
      def get_train_dataset_sizes(
          self, data_param_list, datasets, epoch, shard_epoch=None
      ):
          num_shards = [
              self.get_split_num_data_shards(param["split"])[param["key"]]
              for param in data_param_list
          ]
          data_sizes = []
          for (key, d), num_shard in zip(datasets, num_shards):
              my_data_sizes = self._training_data_sizes[key]
              shard_ind = self.get_shard_id(num_shard, epoch, shard_epoch)
              if shard_ind not in my_data_sizes:
                  my_data_sizes[shard_ind] = len(d)
              known_size = max(my_data_sizes.values())
              data_sizes.append(
                  # If we don't know the data size of the shard yet,
                  # use the the max known data size to approximate.
                  # Note that we preprocess shards by a designated shard size
                  # and put any remaining data at the end into the last shard so
                  # the max shard size approximation is almost correct before loading
                  # the last shard; after loading the last shard, it will have the
                  # exact data sizes of the whole data size.
                  (key, sum(my_data_sizes.get(i, known_size) for i in range(num_shard)))
              )
          logger.info(
              f"estimated total data sizes of all shards used in sampling ratios: {data_sizes}. "
              "Note that if the data a shard has not been loaded yet, use the max known data size to approximate"
          )
          return [s for _, s in data_sizes]
  
      def get_train_sampling_ratios(
          self, data_param_list, datasets, epoch=1, shard_epoch=None
      ):
          data_sizes = self.get_train_dataset_sizes(
              data_param_list, datasets, epoch, shard_epoch
          )
          sampling_func = self.sampling_method.sampling_method_selector()
          sample_ratios = sampling_func(data_sizes) if sampling_func is not None else None
          return sample_ratios
  
      def get_sampling_ratios(self, data_param_list, datasets, epoch, shard_epoch=None):
          if self.args.sampling_weights_from_file:
              weights = load_sampling_weights(self.args.sampling_weights_from_file)
              sample_ratios = [weights[k] for k, _ in datasets]
              logger.info(
                  "| ignoring --sampling-weights when loadding sampling weights "
                  f"from file {self.args.sampling_weights_from_file}"
              )
          elif self.args.sampling_weights:
              sample_ratios = [self.args.sampling_weights[k] for k, _ in datasets]
          else:
              sample_ratios = self.get_train_sampling_ratios(
                  data_param_list, datasets, epoch, shard_epoch
              )
  
          if sample_ratios is not None:
              logger.info(
                  "| Upsample ratios: {}".format(
                      list(zip(map(lambda x: x["key"], data_param_list), sample_ratios))
                  )
              )
              assert len(sample_ratios) == len(datasets)
          return sample_ratios
  
      def load_split_datasets(
          self, split, training, epoch=1, combine=False, shard_epoch=None, **kwargs
      ):
          data_param_list = self.get_split_data_param_list(
              split, epoch, shard_epoch=shard_epoch
          )
          langpairs_sharing_datasets = (
              {} if self.args.enable_reservsed_directions_shared_datasets else None
          )
          datasets = [
              (
                  param["key"],
                  self.load_a_dataset(
                      combine=combine,
                      langpairs_sharing_datasets=langpairs_sharing_datasets,
                      **param,
                  ),
              )
              for param in data_param_list
          ]
          return datasets, data_param_list
  
      def load_into_concat_dataset(self, split, datasets, data_param_list):
          if self.args.lang_tok_replacing_bos_eos:
              # TODO: to investigate why TransformEosLangPairDataset doesn't work with ConcatDataset
              return SampledMultiDataset(
                  OrderedDict(datasets),
                  sampling_ratios=None,
                  eval_key=None,
                  collate_format=CollateFormat.single,
                  virtual_size=None,
                  split=split,
              )
          return ConcatDataset([d for _, d in datasets])
  
      def load_sampled_multi_epoch_dataset(
          self, split, training, epoch=0, combine=False, shard_epoch=None, **kwargs
      ):
          datasets, data_param_list = self.load_split_datasets(
              split, training, epoch, combine, shard_epoch=shard_epoch, **kwargs
          )
          if training and split == getattr(self.args, "train_subset", None):
              sample_ratios = self.get_sampling_ratios(data_param_list, datasets, epoch)
              return SampledMultiEpochDataset(
                  OrderedDict(datasets),
                  epoch=epoch,
                  shard_epoch=shard_epoch,
                  # valid and test datasets will be degenerate to concating datasets:
                  sampling_ratios=sample_ratios,
                  eval_key=None,
                  collate_format=CollateFormat.single,
                  virtual_size=self.args.virtual_data_size,
                  split=split,
                  virtual_epoch_size=self.args.virtual_epoch_size,
                  # if not using lang_tok altering, simplified to use the same collater
                  shared_collater=self._shared_collater(),
              )
          else:
              return self.load_into_concat_dataset(split, datasets, data_param_list)
  