def _multiples_and_weights(config):
  """Helper for weighting GLUE datasets.

  Concatenating all the train sets together and then shuffling the examples
  causes large datasets to dominate the training, resulting in poor performance
  on small datasets. This has some hacky logic to produce (1) "multiples" for
  each dataset so the multi-task train set contains small datasets multiple
  times, so those examples are seen more often and (2) weights for each dataset,
  which also allows for smaller datasets to have influence on training.
  Overall the effect is that tasks are weighted according to
  dataset_size^config.task_weight_exponent.

  Args:
    config: a configure.Config object
  Returns:
    How many copies and weights for each dataset.
  """

  dataset_sizes = {
      "cola": 8551,
      "mnli": 392702,
      "mrpc": 7336,
      "qnli": 108436,
      "qqp": 363869,
      "sst": 67349,
      "sts": 11498,
      "rte": 2490
  }

  def map_values(f, d):
    return {k: f(v) for k, v in d.items()}

  def map_kv(f, d):
    return {k: f(k, v) for k, v in d.items()}

  def normalize(d):
    total = float(sum(d.values()))
    return map_values(lambda v: v / total, d)

  dataset_weights = map_values(lambda s: s ** config.task_weight_exponent,
                               dataset_sizes)
  dataset_weights = normalize(dataset_weights)
  correction = dataset_sizes["mnli"] / dataset_weights["mnli"]
  dataset_tgts = map_values(lambda v: v * correction, dataset_weights)
  dataset_multiples = map_kv(
      lambda task, tgt: round((tgt + 0.01) / dataset_sizes[task]), dataset_tgts)
  new_dataset_sizes = map_kv(
      lambda task, multiple: dataset_sizes[task] * multiple, dataset_multiples)
  weights_after_multiples = map_values(
      lambda v: v * len(dataset_sizes),
      normalize({task: dataset_weights[task] / new_dataset_sizes[task]
                 for task in new_dataset_sizes}))

  return dataset_multiples, weights_after_multiples