Transform cleanup (#176)

* Preference swap clean

* Split is a transform

* Seperate transform module

* Revert accidentl file push

* Fix tests

Author: Jacob-Chmura

aif_gen/cli/commands/split.py

@@ -4,7 +4,7 @@
 
 import click
 
-import aif_gen.dataset.split.functional as F
+import aif_gen.transforms.functional as F
 from aif_gen.dataset.continual_alignment_dataset import (
     ContinualAlignmentDataset,
 )
@@ -73,15 +73,13 @@ def split(
 
     seed_everything(random_seed)
     logging.info(f'Splitting dataset with test_sample_ratio={test_sample_ratio}')
-    dataset = F.split(dataset, test_ratio=test_sample_ratio)
+    transformed_dataset = F.split_transform(dataset, test_ratio=test_sample_ratio)
     logging.info(f'Writing dataset to: {output_file}')
-    dataset.to_json(output_file)
-    logging.info(f'Wrote {dataset.num_samples} samples to: {output_file}')
+    transformed_dataset.to_json(output_file)
+    logging.info(f'Wrote {transformed_dataset.num_samples} samples to: {output_file}')
 
     if hf_repo_id_out is not None:
         upload_to_hf(hf_repo_id_out, output_file)
         logging.info(f'Uploaded dataset to HuggingFace repo: {hf_repo_id_out}')
     else:
         logging.info(f'No HuggingFace repo specified for upload.')
-
-    return

aif_gen/cli/commands/transform.py

@@ -4,7 +4,7 @@
 
 import click
 
-import aif_gen.dataset.transforms.functional as F
+import aif_gen.transforms.functional as F
 from aif_gen.dataset.continual_alignment_dataset import (
     ContinualAlignmentDataset,
 )

aif_gen/dataset/split/__init__.py

@@ -0,0 +1,4 @@
+from aif_gen.transforms.base import DatasetTransform
+from aif_gen.transforms.preference_swap_transform import PreferenceSwapTransform
+from aif_gen.transforms.split_transform import SplitTransform
+from aif_gen.transforms.functional import *

aif_gen/dataset/split/functional.py

@@ -1,26 +1,19 @@
 from abc import ABC, abstractmethod
-from typing import Any, Union
+from typing import Any
 
-from aif_gen.dataset import AlignmentDataset, ContinualAlignmentDataset
-
-# Typedef for convenience
-Dataset = Union[ContinualAlignmentDataset, AlignmentDataset]
+from aif_gen.typing import Dataset
 
 
 class DatasetTransform(ABC):
     r"""Base class for transforming Alignment Datasets."""
 
     @abstractmethod
-    def apply(
-        self, dataset: Dataset, in_place: bool = False, *args: Any, **kwargs: Any
-    ) -> Dataset:
+    def apply(self, dataset: Dataset, in_place: bool = False) -> Dataset:
         r"""Apply the transform onto a dataset.
 
         Args:
             dataset (Union[ContinualAlignmentDataset, AlignmentDataset]): The dataset to transform.
             in_place: Whether to apply the transform in-place or return a new dataset.
-            args (Any): Optional positional arguments.
-            kwargs (Any): Optional keyword arguments.
 
         Returns:
             Union[ContinualAlignmentDataset, AlignmentDataset]: The transformed dataset.
@@ -30,8 +23,4 @@ def __call__(self, dataset: Dataset, *args: Any, **kwargs: Any) -> Dataset:
         return self.apply(dataset, *args, **kwargs)
 
     def __str__(self) -> str:
-        r"""Returns the type of Dataset transform."""
         return self.__class__.__name__
-
-    def _is_dataset_continual(self, dataset: Dataset) -> bool:
-        return isinstance(dataset, ContinualAlignmentDataset)

aif_gen/dataset/transforms/__init__.py

@@ -1,5 +1,6 @@
 from .base import Dataset
 from .preference_swap_transform import PreferenceSwapTransform
+from .split_transform import SplitTransform
 
 
 def preference_swap_transform(
@@ -20,3 +21,23 @@ def preference_swap_transform(
     """
     transform = PreferenceSwapTransform(swap_probability)
     return transform(dataset, in_place=in_place)
+
+
+def split_transform(
+    dataset: Dataset, test_ratio: float, in_place: bool = False
+) -> Dataset:
+    r"""Splits a Dataset training data into train and test datasets.
+
+    Args:
+        dataset (Union[ContinualAlignmentDataset, AlignmentDataset]): The dataset to transform.
+        in_place: Whether to apply the transform in-place or return a new dataset.
+        test_ratio (float): The test ratio to split the dataset with.
+
+    Returns:
+        Union[ContinualAlignmentDataset, AlignmentDataset]: The transformed dataset.
+
+    Raises:
+        ValueError: If a dataset in the Continual Dataset has test data.
+    """
+    transform = SplitTransform(test_ratio)
+    return transform(dataset, in_place=in_place)

aif_gen/transforms/__init__.py

@@ -49,9 +49,7 @@ def apply(self, dataset: Dataset, in_place: bool = False) -> Dataset:
         if self.swap_probability == 0:
             return dataset if in_place else copy.deepcopy(dataset)
 
-        if self._is_dataset_continual(dataset):
-            # This assert is here to make mypy happy
-            assert isinstance(dataset, ContinualAlignmentDataset)
+        if isinstance(dataset, ContinualAlignmentDataset):
             if in_place:
                 for i in range(dataset.num_datasets):
                     dataset.datasets[i] = self._apply(dataset.datasets[i], in_place)
@@ -79,10 +77,10 @@ def _apply_inplace(
     ) -> AlignmentDataset:
         for i in range(len(dataset)):
             if swap_outcomes[i]:
-                chosen = dataset.samples[i].chosen
-                rejected = dataset.samples[i].rejected
-                dataset.samples[i].chosen = rejected
-                dataset.samples[i].rejected = chosen
+                dataset.samples[i].chosen, dataset.samples[i].rejected = (
+                    dataset.samples[i].rejected,
+                    dataset.samples[i].chosen,
+                )
         return dataset
 
     def _apply_copy(
@@ -103,8 +101,6 @@ def _apply_copy(
             train_frac=dataset.train_frac,
         )
 
-    def _validate_swap_probability(self, swap_probability: float) -> None:
-        if not 0 <= swap_probability <= 1:
-            raise ValueError(
-                f'Expected a swap probability in the range [0, 1] but got: {swap_probability}'
-            )
+    def _validate_swap_probability(self, swap_prob: float) -> None:
+        if not 0 <= swap_prob <= 1:
+            raise ValueError(f'Swap probability must be in [0, 1], got: {swap_prob}')

aif_gen/dataset/transforms/base.py aif_gen/transforms/base.pyaif_gen/dataset/transforms/base.py renamed to aif_gen/transforms/base.py

@@ -0,0 +1,87 @@
+from aif_gen.dataset import AlignmentDataset, ContinualAlignmentDataset
+
+from .base import Dataset, DatasetTransform
+
+
+class SplitTransform(DatasetTransform):
+    r"""SplitTransform splits the training data into train/test datasets.
+
+    Args:
+        dataset (Union[ContinualAlignmentDataset, AlignmentDataset]): The dataset to transform.
+        test_ratio (float): The test ratio to split the dataset with.
+
+    Returns:
+        Union[ContinualAlignmentDataset, AlignmentDataset]: The transformed dataset.
+
+    Raises:
+        ValueError: If the test ratio is not in the range [0, 1].
+    """
+
+    def __init__(self, test_ratio: float) -> None:
+        self._validate_test_ratio(test_ratio)
+        self._test_ratio = test_ratio
+
+    @property
+    def test_ratio(self) -> float:
+        r"""float: The test ratio to split the dataset with."""
+        return self._test_ratio
+
+    @test_ratio.setter
+    def test_ratio(self, test_ratio: float) -> None:
+        self._validate_test_ratio(test_ratio)
+        self._test_ratio = test_ratio
+
+    def apply(self, dataset: Dataset, in_place: bool = False) -> Dataset:
+        r"""Splits a ContinualAlignmentDataset's training data into train and test datasets.
+
+        Args:
+            dataset (ContinualAlignmentDataset): The dataset to split.
+            in_place (bool): Whether to apply the transform in-place or return a new dataset.
+
+        Returns:
+            ContinualAlignmentDataset: The dataset with test data included.
+
+        Raises:
+            ValueError: If a dataset in the Continual Dataset has test data.
+        """
+        self._check_test_frac_empty(dataset)
+        if isinstance(dataset, ContinualAlignmentDataset):
+            if in_place:
+                for i in range(dataset.num_datasets):
+                    dataset.datasets[i].train_frac = 1 - self.test_ratio
+                return dataset
+            else:
+                transformed_datasets = []
+                for data in dataset.datasets:
+                    transformed_datasets.append(
+                        AlignmentDataset(
+                            data.task,
+                            data.samples,
+                            train_frac=1 - self.test_ratio,
+                        )
+                    )
+                return ContinualAlignmentDataset(transformed_datasets)
+        else:
+            # This assert is here to make mypy happy
+            assert isinstance(dataset, AlignmentDataset)
+            if in_place:
+                dataset.train_frac = 1 - self.test_ratio
+                return dataset
+            else:
+                return AlignmentDataset(
+                    dataset.task, dataset.samples, train_frac=1 - self.test_ratio
+                )
+
+    def _validate_test_ratio(self, test_ratio: float) -> None:
+        if not 0 <= test_ratio <= 1:
+            raise ValueError(f'Test ratio must be in [0, 1], got: {test_ratio}')
+
+    def _check_test_frac_empty(self, dataset: Dataset) -> None:
+        if isinstance(dataset, ContinualAlignmentDataset):
+            datasets = dataset.datasets
+        else:
+            assert isinstance(dataset, AlignmentDataset)
+            datasets = [dataset]
+        for dataset in datasets:
+            if dataset.test_frac != 0:
+                raise ValueError('AlignmentDataset cannot have test data for splitting')