Update synthesis with usability feedback (#1939)

Author: jgreer013

src/oumi/core/configs/params/synthesis_params.py

@@ -176,19 +176,19 @@ def __post_init__(self):
 
 
 @dataclass
-class PermutableAttributeValue:
-    """Value to be used for the attribute."""
+class SampledAttributeValue:
+    """Value to be sampled for the attribute."""
 
     id: str
     """ID to be used when referencing the attribute value during synthesis."""
 
-    value: str
-    """Value to be used for the attribute.
-    Referenced as {attribute_id.value}"""
+    name: str
+    """Plaintext name of the attribute value.
+    Referenced as {attribute_id}"""
 
     description: str
     """Description of the attribute value.
-    Referenced as {attribute_id.value.description}"""
+    Referenced as {attribute_id.description}"""
 
     sample_rate: Optional[float] = None
     """Sample rate for the attribute value. If not specified, will assume uniform
@@ -197,34 +197,34 @@ class PermutableAttributeValue:
     def __post_init__(self):
         """Verifies/populates params."""
         if not self.id:
-            raise ValueError("PermutableAttributeValue.id cannot be empty.")
-        if not self.value:
-            raise ValueError("PermutableAttributeValue.value cannot be empty.")
+            raise ValueError("SampledAttributeValue.id cannot be empty.")
+        if not self.name:
+            raise ValueError("SampledAttributeValue.name cannot be empty.")
         if not self.description:
-            raise ValueError("PermutableAttributeValue.description cannot be empty.")
+            raise ValueError("SampledAttributeValue.description cannot be empty.")
         if self.sample_rate is not None and (
             self.sample_rate < 0 or self.sample_rate > 1
         ):
             raise ValueError(
-                "PermutableAttributeValue.sample_rate must be between 0 and 1."
+                "SampledAttributeValue.sample_rate must be between 0 and 1."
             )
 
 
 @dataclass
-class PermutableAttribute:
-    """Attributes to be varied across the dataset."""
+class SampledAttribute:
+    """Attributes to be sampled across the dataset."""
 
     id: str
     """ID to be used when referencing the attribute during synthesis."""
 
-    attribute: str
-    """Plaintext name of the attribute. Referenced as {attribute_id}"""
+    name: str
+    """Plaintext name of the attribute. Referenced as {id.parent}"""
 
     description: str
-    """Description of the attribute. Referenced as {attribute_id.description}"""
+    """Description of the attribute. Referenced as {id.parent.description}"""
 
-    possible_values: list[PermutableAttributeValue]
-    """Type of the attribute."""
+    possible_values: list[SampledAttributeValue]
+    """Values to be sampled for the attribute."""
 
     def get_value_distribution(self) -> dict[str, float]:
         """Get the distribution of attribute values."""
@@ -236,13 +236,13 @@ def get_value_distribution(self) -> dict[str, float]:
     def __post_init__(self):
         """Verifies/populates params."""
         if not self.id:
-            raise ValueError("PermutableAttribute.id cannot be empty.")
-        if not self.attribute:
-            raise ValueError("PermutableAttribute.attribute cannot be empty.")
+            raise ValueError("SampledAttribute.id cannot be empty.")
+        if not self.name:
+            raise ValueError("SampledAttribute.name cannot be empty.")
         if not self.description:
-            raise ValueError("PermutableAttribute.description cannot be empty.")
+            raise ValueError("SampledAttribute.description cannot be empty.")
         if not self.possible_values:
-            raise ValueError("PermutableAttribute.possible_values cannot be empty.")
+            raise ValueError("SampledAttribute.possible_values cannot be empty.")
 
         value_ids = []
         sample_rates = []
@@ -252,9 +252,7 @@ def __post_init__(self):
 
         value_ids_set = set(value_ids)
         if len(value_ids) != len(value_ids_set):
-            raise ValueError(
-                "PermutableAttribute.possible_values must have unique IDs."
-            )
+            raise ValueError("SampledAttribute.possible_values must have unique IDs.")
 
         # Normalize sample rates
         normalized_sample_rates = []
@@ -267,7 +265,7 @@ def __post_init__(self):
                 undefined_sample_rate_count += 1
 
             if defined_sample_rate > 1.0:
-                raise ValueError("PermutableAttribute.possible_values must sum to 1.0.")
+                raise ValueError("SampledAttribute.possible_values must sum to 1.0.")
 
         # Assign remaining sample rate to undefined sample rates
         remaining_sample_rate = 1.0 - defined_sample_rate
@@ -517,7 +515,7 @@ class GeneralSynthesisParams(BaseParams):
     Examples will be enumerated during sampling, and attributes can be referenced as
     attributes when generating new attributes."""
 
-    permutable_attributes: Optional[list[PermutableAttribute]] = None
+    sampled_attributes: Optional[list[SampledAttribute]] = None
     """Attributes to be varied across the dataset.
 
     Attributes each have a set of possible values which will be randomly sampled
@@ -636,18 +634,18 @@ def _check_example_source_attribute_ids(self, all_attribute_ids: set[str]) -> No
             for new_key in example_keys:
                 self._check_attribute_ids(all_attribute_ids, new_key)
 
-    def _check_permutable_attribute_ids(self, all_attribute_ids: set[str]) -> None:
-        """Check attribute IDs from permutable attributes for uniqueness."""
-        if self.permutable_attributes is None:
+    def _check_sampled_attribute_ids(self, all_attribute_ids: set[str]) -> None:
+        """Check attribute IDs from sampled attributes for uniqueness."""
+        if self.sampled_attributes is None:
             return
 
-        if len(self.permutable_attributes) == 0:
+        if len(self.sampled_attributes) == 0:
             raise ValueError(
-                "GeneralSynthesisParams.permutable_attributes cannot be empty."
+                "GeneralSynthesisParams.sampled_attributes cannot be empty."
             )
 
-        for permutable_attribute in self.permutable_attributes:
-            attribute_id = permutable_attribute.id
+        for sampled_attribute in self.sampled_attributes:
+            attribute_id = sampled_attribute.id
             self._check_attribute_ids(all_attribute_ids, attribute_id)
 
     def _check_generated_attribute_ids(self, all_attribute_ids: set[str]) -> None:
@@ -716,7 +714,7 @@ def __post_init__(self):
         self._check_dataset_source_attribute_ids(all_attribute_ids)
         self._check_document_source_attribute_ids(all_attribute_ids)
         self._check_example_source_attribute_ids(all_attribute_ids)
-        self._check_permutable_attribute_ids(all_attribute_ids)
+        self._check_sampled_attribute_ids(all_attribute_ids)
         self._check_generated_attribute_ids(all_attribute_ids)
         self._check_transformed_attribute_ids(all_attribute_ids)
         self._check_passthrough_attribute_ids()

src/oumi/core/synthesis/attribute_formatter.py

@@ -16,40 +16,52 @@
 from oumi.utils.placeholders import resolve_placeholders
 
 
-class _AttributeValueInfo:
-    """Information about a value of a permutable attribute.
+class _AttributeParentInfo:
+    """Information about a parent of a sampled attribute."""
 
-    Used to format the string for a sample.
-    """
-
-    def __init__(self, value_name: str, value_description: str):
-        """Initialize the attribute value info."""
-        self._value_name = value_name
-        self.description = value_description
+    def __init__(self, parent_name: str, parent_description: str):
+        """Initialize the attribute parent info."""
+        self._parent_name = parent_name
+        self.description = parent_description
 
     def __str__(self) -> str:
-        return self._value_name
+        return self._parent_name
 
 
 class _AttributeInfo:
-    """Information about a permutable attribute.
+    """Information about a sampled attribute.
 
     Used to format the string for a sample.
+
+    Example:
+        attribute_id: "complexity"
+        parent_name: "Complexity"
+        parent_description: "The complexity of the text."
+        value_name: "High"
+        value_description: "The text is complex."
+
+    Formatting string:
+        {complexity.parent} ({complexity.parent.description})
+        {complexity} ({complexity.description})
+
+    Result:
+        Complexity (The complexity of the text.)
+        High (The text is complex.)
     """
 
     def __init__(
         self,
         attribute_id: str,
-        attribute_name: str,
-        attribute_description: str,
+        parent_name: str,
+        parent_description: str,
         value_name: str,
         value_description: str,
     ):
         """Initialize the attribute value info."""
         self.attribute_id = attribute_id
-        self._attribute_name = attribute_name
-        self.description = attribute_description
-        self.value = _AttributeValueInfo(value_name, value_description)
+        self._attribute_name = value_name
+        self.description = value_description
+        self.parent = _AttributeParentInfo(parent_name, parent_description)
 
     def __str__(self) -> str:
         return self._attribute_name
@@ -65,22 +77,25 @@ class AttributeFormatter:
     def __init__(self, params: GeneralSynthesisParams):
         """Initialize the formatter."""
         self._params = params
-        self._permutable_attribute_map = (
-            {perm_attr.id: perm_attr for perm_attr in params.permutable_attributes}
-            if params.permutable_attributes
+        self._sampled_attribute_map = (
+            {
+                sampled_attr.id: sampled_attr
+                for sampled_attr in params.sampled_attributes
+            }
+            if params.sampled_attributes
             else {}
         )
-        self._permutable_attribute_info = {}
+        self._sampled_attribute_info = {}
 
         # Pre-compute the attribute info for each possible value
-        for attribute_id, attribute in self._permutable_attribute_map.items():
+        for attribute_id, attribute in self._sampled_attribute_map.items():
             for value in attribute.possible_values:
                 key = (attribute_id, value.id)
-                self._permutable_attribute_info[key] = _AttributeInfo(
+                self._sampled_attribute_info[key] = _AttributeInfo(
                     attribute_id=attribute_id,
-                    attribute_name=attribute.attribute,
-                    attribute_description=attribute.description,
-                    value_name=value.value,
+                    parent_name=attribute.name,
+                    parent_description=attribute.description,
+                    value_name=value.name,
                     value_description=value.description,
                 )
 
@@ -102,9 +117,9 @@ def format(
         """
         attr_values = {}
         for attribute_id, attribute_value in sample.items():
-            if self._is_permutable_attribute(attribute_id):
+            if self._is_sampled_attribute(attribute_id):
                 value_id = attribute_value
-                attr_values[attribute_id] = self._get_permutable_attribute_value_info(
+                attr_values[attribute_id] = self._get_sampled_attribute_value_info(
                     attribute_id, value_id
                 )
             else:
@@ -117,17 +132,17 @@ def format(
         )
         return formatted_string
 
-    def _is_permutable_attribute(self, attribute_id: str) -> bool:
-        """Check if the attribute is a permutable attribute."""
-        return attribute_id in self._permutable_attribute_map
+    def _is_sampled_attribute(self, attribute_id: str) -> bool:
+        """Check if the attribute is a sampled attribute."""
+        return attribute_id in self._sampled_attribute_map
 
-    def _get_permutable_attribute_value_info(
+    def _get_sampled_attribute_value_info(
         self, attribute_id: str, attribute_value_id: str
     ) -> _AttributeInfo:
-        """Get the string representation information for a permutable attribute."""
+        """Get the string representation information for a sampled attribute."""
         key = (attribute_id, attribute_value_id)
-        if key in self._permutable_attribute_info:
-            return self._permutable_attribute_info[key]
+        if key in self._sampled_attribute_info:
+            return self._sampled_attribute_info[key]
 
         raise ValueError(
             f"Attribute value {attribute_value_id} not found for "

src/oumi/core/synthesis/dataset_planner.py

@@ -21,7 +21,7 @@
     DocumentSource,
     ExampleSource,
     GeneralSynthesisParams,
-    PermutableAttribute,
+    SampledAttribute,
 )
 from oumi.core.synthesis.dataset_ingestion import DatasetReader
 from oumi.core.synthesis.document_ingestion import DocumentReader, DocumentSegmenter
@@ -84,7 +84,7 @@ def plan(
         )
 
         permutable_attribute_samples = self._plan_permutable_attributes(
-            synthesis_params.permutable_attributes,
+            synthesis_params.sampled_attributes,
             synthesis_params.combination_sampling,
             sample_count,
         )
@@ -190,7 +190,7 @@ def _ingest_document_sources(
 
     def _plan_permutable_attributes(
         self,
-        permutable_attributes: Optional[list[PermutableAttribute]],
+        permutable_attributes: Optional[list[SampledAttribute]],
         combination_sampling: Optional[list[AttributeCombination]],
         sample_count: int,
     ) -> list[dict]: