def load_multi_dataset(
    datasets_config: List[Dict[str, Union[int, str, Sequence[str]]]],
    name_to_features: Dict[str, Any],
    preprocess_fn: Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]],
    collater_fn: Callable[[Dict[str, tf.Tensor]], Dict[str, tf.Tensor]],
    is_training: bool,
    per_device_batch_size: int,
    local_device_count: int,
    host_count: int,
    host_id: int,
    seed: int = 0,
    pad_eval: bool = False,
) -> tf.data.Dataset:
  """Create TF dataset from different data sources.

  Args:
    datasets_config: a list of datasets to load.
    name_to_features: features to decode from a record in binary format.
    preprocess_fn: performs preprocessing on examples.
    collater_fn: performs preprocessing on the whole batch.
    is_training: if true, shuffles
    per_device_batch_size: total batch size per host.
    local_device_count: number of devices, used for batching.
    host_count: number of hosts, used for sharding.
    host_id: ID of current host, used for sharding.
    seed: random seed for non-deterministic operations.
    pad_eval: if true, pad rather than discard last batch for evaluation.

  Returns:
    tf.data.Dataset. Input data for training.
  """
  datasets = []
  for dataset_config in datasets_config:
    decode_fn_per_dataset = make_decode_fn(
        name_to_features=name_to_features,
        samples_per_example=dataset_config['samples_per_example'],
    )
    datasets.append(
        load_dataset(
            patterns=dataset_config['patterns'],
            decode_fn=decode_fn_per_dataset,
            preprocess_fn=preprocess_fn,
            collater_fn=collater_fn,
            is_training=is_training,
            per_device_batch_size=per_device_batch_size,
            local_device_count=local_device_count,
            host_count=host_count,
            host_id=host_id,
            seed=seed,
            pad_eval=pad_eval,
        ))
  if len(datasets) > 1:
    dataset = tf.data.experimental.sample_from_datasets(
        datasets, weights=None, seed=seed)
  else:
    dataset = datasets[0]
  return dataset
Create TF dataset from different data sources.

  Args:
    datasets_config: a list of datasets to load.
    name_to_features: features to decode from a record in binary format.
    preprocess_fn: performs preprocessing on examples.
    collater_fn: performs preprocessing on the whole batch.
    is_training: if true, shuffles
    per_device_batch_size: total batch size per host.
    local_device_count: number of devices, used for batching.
    host_count: number of hosts, used for sharding.
    host_id: ID of current host, used for sharding.
    seed: random seed for non-deterministic operations.
    pad_eval: if true, pad rather than discard last batch for evaluation.

  Returns:
    tf.data.Dataset. Input data for training.
  
language/mentionmemory/utils/data_utils.py