def generate_data(self,
                    situation_specifications,
                    max_examples=None,
                    other_objects_sample_percentage=0.5,
                    split_type="uniform",
                    visualize_per_template=0,
                    visualize_per_split=0,
                    train_percentage=0.8,
                    min_other_objects=0,
                    max_objects=12,
                    cut_off_target_length=25,
                    max_distract_objects=1,
                    max_examples_per_target=None):
    """Generate data pairs given situations and assign them to corresponding splits."""
    example_count = 0
    dropped_examples = 0
    dropped_situations = 0
    for t_num, template_derivations in self._grammar.all_derivations.items():
      visualized_per_template = 0
      visualized_per_split = {split: 0 for split in self._possible_splits}
      for derivation_num, derivation in enumerate(template_derivations):
        arguments = []
        derivation.meaning(arguments)
        assert len(arguments) == 1, "Only one target object supported."
        adverb = ""
        for word in derivation.words():
          if word in self._vocabulary.get_adverbs():
            adverb = word
        _, target_predicate, ref_predicate = arguments.pop().to_predicate(
            return_ref_predicate=True)
        possible_target_objects = self.generate_possible_targets(
            referred_size=self._vocabulary.translate_word(
                target_predicate["size"]),
            referred_color=self._vocabulary.translate_word(
                target_predicate["color"]),
            referred_shape=self._vocabulary.translate_word(
                target_predicate["noun"]))
        for target_size, target_color, target_shape in possible_target_objects:
          relevant_situations = situation_specifications[target_shape][
              target_color][target_size]
          relevant_situations = self.remove_invalid_siutations(
              relevant_situations, target_predicate, ref_predicate)
          num_relevant_situations = len(relevant_situations)
          if num_relevant_situations == 0:
            continue
          if max_examples_per_target:
            num_relevant_situations = min(max_examples_per_target,
                                          num_relevant_situations)
            relevant_situations = random.sample(
                relevant_situations, k=num_relevant_situations)
          idx_to_visualize = random.sample(
              list(range(num_relevant_situations)), k=1).pop()

          if split_type == "uniform":
            idx_for_train = random.sample(
                list(range(num_relevant_situations)),
                k=int(num_relevant_situations * train_percentage))
            idx_for_train = set(idx_for_train)
          for i, relevant_situation in enumerate(relevant_situations):
            visualize = False
            if (example_count + 1) % 10000 == 0:
              logging.info(
                  "%02d/%02d templates, %05d/%05d derivations, %08d examples",
                  t_num, len(self._grammar.all_derivations), derivation_num,
                  len(template_derivations), example_count + 1)
            valid_init = self.initialize_world_from_spec(
                relevant_situation,
                referred_size=target_predicate["size"],
                referred_color=target_predicate["color"],
                referred_shape=target_predicate["noun"],
                actual_size=target_size,
                sample_percentage=other_objects_sample_percentage,
                min_other_objects=min_other_objects,
                max_objects=max_objects,
                ref_predicate=ref_predicate,
                max_distract_objects=max_distract_objects)
            if not valid_init:
              dropped_situations += 1
              if i == idx_to_visualize:
                idx_to_visualize += 1
              continue
            situation = self._world.get_current_situation()
            assert situation.direction_to_target == relevant_situation[
                "direction_to_target"]
            assert situation.distance_to_target == relevant_situation[
                "distance_to_target"]
            target_commands, _, target_action = self.demonstrate_command(
                derivation, initial_situation=situation)

            if i == idx_to_visualize:
              visualize = True
            if visualized_per_template >= visualize_per_template:
              visualize = False
            if adverb and visualized_per_template <= visualize_per_template:
              visualize = True
            if split_type == "uniform":
              if i in idx_for_train:
                splits = ["train"]
              else:
                splits = ["test"]
            elif split_type == "generalization":
              splits = self.assign_splits(
                  target_size, target_color, target_shape, target_action,
                  situation.direction_to_target, target_predicate,
                  self._vocabulary.translate_word(adverb))
              if not splits:
                splits = ["train"]
              elif len(splits) > 1:
                dropped_examples += 1
                self._world.clear_situation()
                continue
            elif split_type == "target_lengths":
              if len(target_commands) > cut_off_target_length:
                splits = ["test"]
              else:
                splits = ["train"]
            elif split_type == "spatial_relation":
              splits = self.assign_relation_splits(target_color, target_shape,
                                                   ref_predicate,
                                                   self._world.reference_object)
              if not splits:
                splits = ["train"]
              elif len(splits) > 1:
                dropped_examples += 1
                self._world.clear_situation()
                continue
            else:
              raise ValueError("Unknown split_type in .get_data_pairs().")
            if visualized_per_split[splits[0]] <= visualize_per_split:
              visualized_per_split[splits[0]] += 1
              visualize = True
            self.fill_example(
                command=derivation.words(),
                derivation=derivation,
                situation=situation,
                target_commands=target_commands,
                verb_in_command=target_action,
                location_prep=ref_predicate["location"],
                target_predicate=target_predicate,
                visualize=visualize,
                adverb=adverb,
                splits=splits)
            for split in splits:
              self._template_identifiers[split].append(t_num)
            example_count += 1
            if visualize:
              visualized_per_template += 1
            self._world.clear_situation()
          if max_examples and example_count >= max_examples:
            return example_count, dropped_examples, dropped_situations
    return example_count, dropped_examples, dropped_situations
Generate data pairs given situations and assign them to corresponding splits.
language/gscan/data/dataset.py