Implement in_place transform and tests

Author: Jacob-Chmura

aif_gen/dataset/transforms/base.py

@@ -26,10 +26,8 @@ def apply(
             Union[ContinualAlignmentDataset, AlignmentDataset]: The transformed dataset.
         """
 
-    def __call__(
-        self, dataset: Dataset, in_place: bool = False, *args: Any, **kwargs: Any
-    ) -> Dataset:
-        return self.apply(dataset, in_place, *args, **kwargs)
+    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."""

aif_gen/dataset/transforms/functional.py

@@ -19,8 +19,4 @@ def preference_swap_transform(
         ValueError: If the swap probability is not in the range [0, 1].
     """
     transform = PreferenceSwapTransform(swap_probability)
-    if in_place:
-        transform(dataset)
-    else:
-        dataset = transform(dataset)
-    return dataset
+    return transform(dataset, in_place=in_place)

aif_gen/dataset/transforms/preference_swap_transform.py

@@ -52,18 +52,42 @@ def apply(self, dataset: Dataset, in_place: bool = False) -> Dataset:
         if self._is_dataset_continual(dataset):
             # This assert is here to make mypy happy
             assert isinstance(dataset, ContinualAlignmentDataset)
-            return ContinualAlignmentDataset(
-                [self._apply(data) for data in dataset.datasets]
-            )
+            if in_place:
+                for i in range(dataset.num_datasets):
+                    dataset.datasets[i] = self._apply(dataset.datasets[i], in_place)
+                return dataset
+            else:
+                return ContinualAlignmentDataset(
+                    [self._apply(data, in_place) for data in dataset.datasets]
+                )
         else:
             # This assert is here to make mypy happy
             assert isinstance(dataset, AlignmentDataset)
-            return self._apply(dataset)
+            return self._apply(dataset, in_place)
 
-    def _apply(self, dataset: AlignmentDataset) -> AlignmentDataset:
+    def _apply(self, dataset: AlignmentDataset, in_place: bool) -> AlignmentDataset:
         swap_outcomes = np.random.binomial(
             n=1, p=self.swap_probability, size=len(dataset)
         )
+        if in_place:
+            return self._apply_inplace(dataset, swap_outcomes)
+        else:
+            return self._apply_copy(dataset, swap_outcomes)
+
+    def _apply_inplace(
+        self, dataset: AlignmentDataset, swap_outcomes: np.ndarray
+    ) -> 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
+        return dataset
+
+    def _apply_copy(
+        self, dataset: AlignmentDataset, swap_outcomes: np.ndarray
+    ) -> AlignmentDataset:
         transformed_samples = []
         for i, sample in enumerate(dataset.samples):
             if swap_outcomes[i]:

test/test_transforms/conftest.py

@@ -0,0 +1,19 @@
+import pytest
+
+from aif_gen.util.seed import seed_everything
+
+
+@pytest.fixture(autouse=True)
+def run_seed_before_tests():
+    seed_everything(1)
+    yield
+
+
+@pytest.fixture(params=[True, False])
+def in_place(request):
+    return request.param
+
+
+@pytest.fixture(params=['call', 'apply', 'functional'])
+def application_type(request):
+    return request.param