def generate_samples(self, data_dir, tmp_dir, dataset_split):
    """Generate just the first 6k samples for training."""
    # If not training, do the same as before.
    if dataset_split != problem.DatasetSplit.TRAIN:
      for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
          data_dir, tmp_dir, dataset_split):
        yield x
      raise StopIteration
    # Now we assume we're training.
    counter = 0
    # The size of this data-set in total is around 614K, we want to sample so
    # that in expectation we take the requested number of samples in 1 go.
    sample_prob = self.how_many_examples_to_sample / float(614000)
    # Let's sample.
    for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
        data_dir, tmp_dir, dataset_split):
      if random.random() > sample_prob:
        continue
      counter += 1
      if counter > self.how_many_examples_to_sample:
        raise StopIteration
      yield x
    # We do it again if we don't have enough samples.
    if counter < self.how_many_examples_to_sample:
      for x in super(TranslateEnroWmtMultiSmall64k, self).generate_samples(
          data_dir, tmp_dir, dataset_split):
        if random.random() > sample_prob:
          continue
        counter += 1
        if counter > self.how_many_examples_to_sample:
          raise StopIteration
        yield x
Generate just the first 6k samples for training.
tensor2tensor/data_generators/translate_enro.py