recoglab_dataset.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Code for generating ReCogLab dataset.

This is a dataset of different reasoning modules which can be strung together,
where the goal is to answer a question about a set of entities.

The dataset is generated by creating a set of modules, where each
module is a sequence of blocks. Each block reveals information about the
entities.

The goal of the dataset is to create modules that are long and
complex, while still being understandable by humans.
"""

import collections
from collections.abc import Mapping
import itertools
from typing import Any, Dict, List, Set

import jax
import ml_collections

from recoglab import common_types
from recoglab import utils
from recoglab.modules import comparison
from recoglab.modules import family
from recoglab.modules import filler
from recoglab.modules import recoglab_base
from recoglab.modules import social_network
from recoglab.modules import syllogisms


# Multiple modules form an EEAODatasetExample.
class ReCogLabDatasetExample:
  """Example of ReCogLab dataset."""

  def __init__(self):
    self.all_modules = []
    self.all_blocks = []
    self.answer_block: recoglab_base.ReCogLabQuestionAnswerBlock | None = None
    self.metadata = {}

  def get_accuracy(self, pred: str) -> float:
    """Evaluates the prediction against the answer block.

    Args:
      pred: The prediction to evaluate.

    Raises:
      AttributeError: answer_block has not been set yet

    Returns:
      Score on question/answer block.
    """
    if not self.answer_block:
      raise AttributeError('Answer block was not defined.')
    return self.answer_block.get_accuracy(pred)

  def add_module(self, module: recoglab_base.ReCogLabModule):
    """Adds a module to the example's list of blocks.

    Args:
      module: The module to add to the example.
    """
    self.all_modules.append(module)
    self.all_blocks += module.get_blocks()
    self.answer_block = module.get_answer_block()
    # We assume that most modules will have empty metadata.
    self.metadata.update(module.get_metadata())

  def get_all_entities(self) -> Set[common_types.Entity]:
    """Returns all the entities in the example."""
    return set(
        itertools.chain.from_iterable(
            module.get_entities() for module in self.all_modules
        )
    )

  def get_prompts(self, block_separator: str = '\n') -> str:
    """Returns the prompts for the example."""
    block_prompts = (
        block.prompt_separator.join(block.prompt) for block in self.all_blocks
    )
    return block_separator.join(block_prompts)

  def get_question(self) -> str:
    """Returns the question for the example.

    Raises:
      AttributeError: answer_block has not been set yet
    """
    if not self.answer_block:
      raise AttributeError('Answer block was not defined.')
    return self.answer_block.get_question()

  def get_answers(self) -> list[str]:
    """Returns the answers for the example.

    Raises:
      AttributeError: answer_block has not been set yet
    """
    if not self.answer_block:
      raise AttributeError('Answer block was not defined.')
    return self.answer_block.answers

  def get_metadata(self) -> Dict[str, str]:
    return self.metadata


class ReCogLabDatasetGenerator:
  """Generator class for ReCogLab Dataset.

  Attributes:
    cfg: Configuration for dataset to generate.
    module_dict: dict
    jax_key: PRNGKey for generating examples.
    split: ['train', 'val', 'test'] split to generate examples for.
  """

  def __init__(
      self,
      dataset_config: ml_collections.ConfigDict,
      split: str,
      seed: int = 0,
  ):
    self.cfg = dataset_config
    self.module_dict = {}
    self.jax_key = jax.random.PRNGKey(seed)
    self.split = common_types.DatasetSplit(split)

    self.jax_key, cfg_module_key = jax.random.split(self.jax_key)
    # Initializes modules
    for c_name in self.cfg.all_module_names:
      module_cfg = self.cfg[c_name]
      cfg_module_key, subkey = jax.random.split(cfg_module_key)
      if module_cfg.name == 'SocialNetworkModule':
        self.module_dict[module_cfg.name] = (
            social_network.SocialNetworkModuleGenerator(
                module_cfg, subkey, split=self.split
            )
        )
      elif module_cfg.name == 'SyllogismModule':
        self.module_dict[module_cfg.name] = syllogisms.SyllogismModuleGenerator(
            module_cfg, subkey, split=self.split
        )
      elif module_cfg.name == 'ComparisonModule':
        self.module_dict[module_cfg.name] = (
            comparison.ComparisonModuleGenerator(
                module_cfg, subkey, split=self.split
            )
        )
      elif module_cfg.name == 'ComparisonValidModule':
        module_generator = comparison.ComparisonModuleGenerator(
            module_cfg, subkey, split=self.split
        )
        # swap the default module generator for the validity module generator
        module_generator.generate_module = (
            module_generator.generate_validity_module
        )
        self.module_dict[module_cfg.name] = module_generator
      elif module_cfg.name == 'FamilyModule':
        self.module_dict[module_cfg.name] = family.FamilyModuleGenerator(
            module_cfg, subkey, split=self.split
        )
      else:
        raise NotImplementedError(
            f'Need to implement this module type: {module_cfg.name}'
        )

  def generate_examples(
      self,
      num_examples: int = 1,
      metadata_rebalance: bool = False,
      metadata_field_to_rebalance: str = '',
  ) -> List[ReCogLabDatasetExample]:
    """Generates num_examples of recoglab_dataset.ReCogLabDatasetExample.

    Args:
      num_examples: The number of examples to generate.
      metadata_rebalance: If true, attempt to balance examples by discarding
        examples. The first 10% of num_examples generated will be emitted to
        gauge a natural distribution of metadata. Then a heuristic is used to
        balance the emitted data.
      metadata_field_to_rebalance: A string of a field to balance around

    Returns:
      A list of ReCogLabDatasetExamples.
    """
    seen_metadata_count = collections.defaultdict(int)
    emitted_metadata_count = collections.defaultdict(int)
    examples = []
    for example_idx in range(num_examples):
      all_used_entities = []
      example = ReCogLabDatasetExample()
      self.jax_key, example_subkey = jax.random.split(self.jax_key)
      for c_name in self.cfg.all_module_names:
        module_cfg = self.cfg[c_name]

        # Try-catch generating modules (can throw DataGenerationError)
        while True:
          if module_cfg.name not in self.module_dict:
            raise NotImplementedError(
                f'Need to implement this module type: {module_cfg.name}'
            )
          try:
            module = self.module_dict[module_cfg.name].generate_module(
                example_subkey, all_used_entities
            )
            break
          except common_types.DataGenerationError:
            _, example_subkey = jax.random.split(example_subkey)
            continue

        # Add filler to module
        if module_cfg.add_filler:
          if module_cfg.filler_type == 'random_text':
            filler_lines = filler.generate_random_text(
                module_cfg.num_filler_lines, example_subkey
            )
          elif module_cfg.filler_type == 'entity_filler':
            if not module_cfg.entity_type:
              raise ValueError(
                  'entity_type must be set if filler_type is entity_filler'
              )
            filler_lines = filler.generate_entity_filler(
                module_cfg.num_filler_lines,
                list(module.get_entities()),
                module_cfg.entity_type,
                example_subkey,
            )
          else:
            raise NotImplementedError(
                f'Need to implement this filler type: {module_cfg.filler_type}'
            )
          # Add filler lines to module
          if module_cfg.filler_position == 'before':
            module.blocks = filler_lines + module.blocks
          elif module_cfg.filler_position == 'after':
            module.blocks += filler_lines
          elif module_cfg.filler_position == 'interspersed':
            _, example_subkey = jax.random.split(example_subkey)
            module.blocks = utils.interleave(
                filler_lines, module.blocks, example_subkey
            )
          else:
            raise NotImplementedError(
                'Need to implement this filler position:'
                f' {module_cfg.filler_position}'
            )
        example.add_module(module)
        if self.cfg.maintain_entity_uniqueness:
          all_used_entities.extend(module.get_entities())
      example_metadata_field = example.get_metadata().get(
          metadata_field_to_rebalance
      )
      seen_metadata_count[example_metadata_field] += 1
      # Check metadata conditions
      if metadata_rebalance:
        heuristic_should_emit = emit_examples_heuristic(
            example_metadata_field,
            seen_metadata_count,
            emitted_metadata_count,
        )
        should_emit = (
            example_idx < int(0.1 * num_examples) or heuristic_should_emit
        )
      else:
        should_emit = True
      if should_emit:
        emitted_metadata_count[example_metadata_field] += 1
        examples.append(example)
    return examples


def emit_examples_heuristic(
    metadata: Any, seen_counts: Mapping[Any, int],
    emitted_counts: Mapping[Any, int]
) -> bool:
  """Checks if recoglab_example should be emitted given the existing metadata.

  The following heuristic is used to balance metadata:
    1. If emitted_counts[metadata] has < 20 emitted examples, Return True
    2. If emitted_counts[metadata] > 2 * emitted_counts[metadata*]
      where metadata* is the field that contains fields are above 0.01 minimum
      propotion threshold of seen_counts.

  Args:
    metadata: the metadata field value to apply heuristic rebalancing
    seen_counts: mapping between metadata fields to seen counts. This is the
      natural data generating distribution of metadata.
    emitted_counts: mapping between metadata fields and emitted data

  Returns:
    True if heuristic indicates to emit example.
  """
  # If a field doesn't occur naturally at least 1% of the time, then it does
  # not occur frequently enough to be eligible for rebalancing.
  min_p = 0.01
  auto_output_threshold = 20

  if emitted_counts.get(metadata, 0) < auto_output_threshold:
    return True
  total_examples = sum(seen_counts.values())
  threshold_for_inclusion = total_examples * min_p
  smallest_count = min(
      count
      for key, count in emitted_counts.items()
      if seen_counts[key] >= threshold_for_inclusion
  )
  if emitted_counts[metadata] < 2 * smallest_count:
    return True
  return False


def generate_dataset(
    dataset_config, split, seed, num_examples, metadata_rebalance_field=''
) -> List[ReCogLabDatasetExample]:
  """Entry point to generate dataset.

  Args:
    dataset_config: The dataset config to generate.
    split: The split to generate.
    seed: The seed to generate.
    num_examples: The number of examples to generate.
    metadata_rebalance_field: If populated, will rebalance the generation around
      a metadata attribute.

  Returns:
    A list of ReCogLabDatasetExamples.
  """
  if metadata_rebalance_field:
    metadata_rebalance = True
  else:
    metadata_rebalance = False

  # Generate the dataset and examples
  dataset_generator = ReCogLabDatasetGenerator(
      dataset_config, split=split, seed=seed
  )
  examples = dataset_generator.generate_examples(
      num_examples=num_examples,
      metadata_rebalance=metadata_rebalance,
      metadata_field_to_rebalance=metadata_rebalance_field,
  )
  return examples