Repository: oumi-ai/oumi
File: src/oumi/core/synthesis/planner.py (Python)

Code before:
            synthesis_params: The synthesis parameters.
            sample_count: The number of samples to plan.

        Returns:
            A list of dictionaries, each representing a sample of the dataset with
            the attribute values for each attribute.
        """
        permutable_attributes = self._plan_permutable_attributes(
            synthesis_params.permutable_attributes,
            synthesis_params.combination_sampling,
            sample_count,
        )

        dataset_plan = permutable_attributes

        return dataset_plan

    def _plan_permutable_attributes(
        self,
        permutable_attributes: Optional[list[PermutableAttribute]],
        combination_sampling: Optional[list[AttributeCombination]],
        sample_count: int,
    ) -> list[dict]:
        if sample_count < 0:
            raise ValueError("Count must be positive")
        elif (
            sample_count == 0
            or permutable_attributes is None
            or len(permutable_attributes) == 0
        ):
            return []

        sampling_overrides = combination_sampling or []

        cumulative_override_probability = sum(
            [combination.sample_rate for combination in sampling_overrides]
        )

        # If cumulative probability is greater than 1, raise an error
        if cumulative_override_probability > 1.0:
            raise ValueError(
                "Cumulative probability of combination sampling must "
                "be less than or equal to 1.0."
            )

        # Generate sample count permutations of the permutable attributes
        samples = [{} for _ in range(sample_count)]
        attribute_distributions = {
            perm_attr.id: perm_attr.get_value_distribution()
            for perm_attr in permutable_attributes
        }

Review comment (suggestion):
```suggestion
        return self._plan_permutable_attributes(
            synthesis_params.permutable_attributes,
            synthesis_params.combination_sampling,
            sample_count,
        )
```

Code after:
        representing a sample of the dataset with a particular attribute value for
        each attribute.

        - Permutable attributes have their values sampled from a distribution.
        - Combination sampling overrides the distribution for particular attribute value
          combinations.

        The final list of dictionaries will be used to create a dataset.

        Args:
            synthesis_params: The synthesis parameters.
            sample_count: The number of samples to plan.

        Returns:
            A list of dictionaries, each representing a sample of the dataset with
            the attribute values for each attribute.
        """
        permutable_attributes = self._plan_permutable_attributes(
            synthesis_params.permutable_attributes,
            synthesis_params.combination_sampling,
            sample_count,
        )

        dataset_plan = permutable_attributes

        return dataset_plan

    def _plan_permutable_attributes(
        self,
        permutable_attributes: Optional[list[PermutableAttribute]],
        combination_sampling: Optional[list[AttributeCombination]],
        sample_count: int,
    ) -> list[dict]:
        if sample_count < 0:
            raise ValueError("Count must be positive")
        elif (
            sample_count == 0
            or permutable_attributes is None
            or len(permutable_attributes) == 0
        ):
            return []

        sampling_overrides = combination_sampling or []

        cumulative_override_probability = sum(
            [combination.sample_rate for combination in sampling_overrides]
        )

        # If cumulative probability is greater than 1, raise an error
        if cumulative_override_probability > 1.0:
            raise ValueError(