Dataset Cleanup (#174)

* WIP

* Fix tests

* update continual dataset

* Pipe instead of union

* Annotate

Author: Jacob-Chmura

aif_gen/dataset/alignment_dataset.py

@@ -1,15 +1,20 @@
+from __future__ import annotations
+
 import json
 import pathlib
 from dataclasses import asdict
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List
 
 from datasets import Dataset
+from pydantic import Field
+from pydantic.dataclasses import dataclass
 
 from aif_gen.task import AlignmentTask
 
 from .alignment_sample import AlignmentDatasetSample
 
 
+@dataclass(slots=True)
 class AlignmentDataset:
     r"""Container object for an Alignment Dataset.
 
@@ -22,38 +27,14 @@ class AlignmentDataset:
         ValueError: If train_frac is not in the interval [0, 1.0]
     """
 
-    def __init__(
-        self,
-        task: AlignmentTask,
-        samples: List[AlignmentDatasetSample],
-        train_frac: float = 1.0,
-    ) -> None:
-        self._task = task
-        self._samples = samples
-
-        if not (0 <= train_frac <= 1):
-            raise ValueError(f'Train fraction must be in [0, 1] but got: {train_frac}')
-        self._train_frac = train_frac
-
-    @property
-    def task(self) -> AlignmentTask:
-        r"""AlignmentTask: The task associated with the AlignmentDataset."""
-        return self._task
-
-    @property
-    def train_frac(self) -> float:
-        r"""Fraction of samples that belong to the training split."""
-        return self._train_frac
+    task: AlignmentTask = Field(frozen=True)
+    samples: List[AlignmentDatasetSample] = Field(frozen=True)
+    train_frac: float = Field(default=1.0, ge=0, le=1)
 
     @property
     def test_frac(self) -> float:
         r"""Fraction of samples that belong to the testing split."""
-        return 1.0 - self._train_frac
-
-    @property
-    def samples(self) -> List[AlignmentDatasetSample]:
-        r"""List[AlignmentDatasetSample]: The list of samples associated with the AlignmentDataset."""
-        return self._samples
+        return 1.0 - self.train_frac
 
     @property
     def train(self) -> List[AlignmentDatasetSample]:
@@ -84,12 +65,12 @@ def __len__(self) -> int:
         return self.num_samples
 
     def __getitem__(
-        self, key: Union[slice, int]
-    ) -> Union[AlignmentDatasetSample, List[AlignmentDatasetSample]]:
+        self, key: slice | int
+    ) -> AlignmentDatasetSample | List[AlignmentDatasetSample]:
         # Slicing directly on the samples
         return self.samples[key]
 
-    def to_json(self, file_path: Union[str, pathlib.Path]) -> None:
+    def to_json(self, file_path: str | pathlib.Path) -> None:
         r"""Save the AlignmentDataset to a json file.
 
         Note: Uses to_dict() under the hood to get a dictionary representation.
@@ -104,26 +85,17 @@ def to_json(self, file_path: Union[str, pathlib.Path]) -> None:
     def to_dict(self) -> Dict[str, Any]:
         r"""Convert the AlignmentDataset to dictionary represenetation.
 
-        Note: This method is the functional inverse of AlignmentDataset.from_dict().
-
         Returns:
             Dict[str, Any]: The dictionary representation of the AlignmentDataset.
         """
         dataset_dict: Dict[str, Any] = {}
         dataset_dict['task'] = self.task.to_dict()
-        dataset_dict['train'] = []
-        dataset_dict['test'] = []
-
-        for sample in self.train:
-            dataset_dict['train'].append(asdict(sample))
-
-        for sample in self.test:
-            dataset_dict['test'].append(asdict(sample))
-
+        dataset_dict['train'] = [asdict(sample) for sample in self.train]
+        dataset_dict['test'] = [asdict(sample) for sample in self.test]
         return dataset_dict
 
     @classmethod
-    def from_json(cls, file_path: Union[str, pathlib.Path]) -> 'AlignmentDataset':
+    def from_json(cls, file_path: str | pathlib.Path) -> AlignmentDataset:
         r"""Load the AlignmentDataset from a json file.
 
         Note: Uses AlignmentDataset.from_dict() under the hood to parse the representation.
@@ -136,11 +108,10 @@ def from_json(cls, file_path: Union[str, pathlib.Path]) -> 'AlignmentDataset':
         """
         with open(file_path, 'r') as f:
             dataset_dict = json.load(f)
-
         return cls.from_dict(dataset_dict)
 
     @classmethod
-    def from_dict(cls, dataset_dict: Dict[str, Any]) -> 'AlignmentDataset':
+    def from_dict(cls, dataset_dict: Dict[str, Any]) -> AlignmentDataset:
         r"""Construct an AlignmentDataset from dictionary representation.
 
         Note:
@@ -161,14 +132,11 @@ def from_dict(cls, dataset_dict: Dict[str, Any]) -> 'AlignmentDataset':
         task = AlignmentTask.from_dict(dataset_dict['task'])
         samples = []
         for sample in dataset_dict['train']:
-            sample = AlignmentDatasetSample(**sample)
-            samples.append(sample)
-
+            samples.append(AlignmentDatasetSample(**sample))
         num_train_samples = len(samples)
 
         for sample in dataset_dict['test']:
-            sample = AlignmentDatasetSample(**sample)
-            samples.append(sample)
+            samples.append(AlignmentDatasetSample(**sample))
 
         train_frac = num_train_samples / len(samples)
         return cls(task, samples, train_frac)
@@ -177,26 +145,22 @@ def to_hf_compatible(self) -> Dict[str, Dataset]:
         r"""Convert the AlignmentDataset to a dictionary compatible with HuggingFace datasets.
 
         Returns:
-            dict[str, Dataset]: The dictionary compatible with HuggingFace datasets.
+            Dict[str, Dataset]: The dictionary compatible with HuggingFace datasets.
         """
-        dataset_dict: Dict[str, Any] = self.to_dict()
-
         hf_dict: Dict[str, Dataset] = {
             'train': Dataset.from_dict(
                 {
-                    'prompt': [sample['prompt'] for sample in dataset_dict['train']],
-                    'chosen': [sample['chosen'] for sample in dataset_dict['train']],
-                    'rejected': [
-                        sample['rejected'] for sample in dataset_dict['train']
-                    ],
+                    'prompt': [sample.prompt for sample in self.train],
+                    'chosen': [sample.chosen for sample in self.train],
+                    'rejected': [sample.rejected for sample in self.train],
                 },
                 split='train',
             ),
             'test': Dataset.from_dict(
                 {
-                    'prompt': [sample['prompt'] for sample in dataset_dict['test']],
-                    'chosen': [sample['chosen'] for sample in dataset_dict['test']],
-                    'rejected': [sample['rejected'] for sample in dataset_dict['test']],
+                    'prompt': [sample.prompt for sample in self.test],
+                    'chosen': [sample.chosen for sample in self.test],
+                    'rejected': [sample.rejected for sample in self.test],
                 },
                 split='test',
             ),

aif_gen/dataset/alignment_sample.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 
 
 @dataclass

aif_gen/dataset/continual_alignment_dataset.py

@@ -1,28 +1,25 @@
+from __future__ import annotations
+
 import json
 import pathlib
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List
 
 from datasets import Dataset
-
-from aif_gen.dataset.alignment_sample import AlignmentDatasetSample
+from pydantic.dataclasses import dataclass
 
 from .alignment_dataset import AlignmentDataset
+from .alignment_sample import AlignmentDatasetSample
 
 
+@dataclass(slots=True)
 class ContinualAlignmentDataset:
     r"""Container object for a Continual Alignment Dataset.
 
     Args:
         datasets (List[ContinualAlignmentDataset]): Temporal list of AlignmentDatasets constituents.
     """
 
-    def __init__(self, datasets: List[AlignmentDataset]) -> None:
-        self._datasets = datasets
-
-    @property
-    def datasets(self) -> List[AlignmentDataset]:
-        r"""List[AlignmentDataset]: The list of AlignmentDataset constituents."""
-        return self._datasets
+    datasets: List[AlignmentDataset]
 
     @property
     def num_datasets(self) -> int:
@@ -39,16 +36,16 @@ def __len__(self) -> int:
         return self.num_samples
 
     def __getitem__(
-        self, key: Union[slice, int]
-    ) -> Union[AlignmentDatasetSample, List[AlignmentDatasetSample]]:
+        self, key: slice | int
+    ) -> AlignmentDatasetSample | List[AlignmentDatasetSample]:
         # Indexing based on **samples** across datasets (not into datasets themselves)
         all_samples = []  # This should probably be cached
         for dataset in self.datasets:
             all_samples.extend(dataset.samples)
         return all_samples[key]
 
     def append(self, dataset: AlignmentDataset) -> None:
-        r"""Append a single AlignmentDataset to the ConitnualAlignmentDataset.
+        r"""Append a single AlignmentDataset to the ContinualAlignmentDataset.
 
         Args:
             dataset (AlignmentDataset): The new dataset to add.
@@ -64,7 +61,7 @@ def append(self, dataset: AlignmentDataset) -> None:
             )
 
     def extend(self, datasets: List[AlignmentDataset]) -> None:
-        r"""Append multiple AlignmentDataset's to the ConitnualAlignmentDataset.
+        r"""Append multiple AlignmentDataset's to the ContinualAlignmentDataset.
 
         Args:
             datasets (List[AlignmentDataset]): The new datasets to add.
@@ -75,7 +72,7 @@ def extend(self, datasets: List[AlignmentDataset]) -> None:
         for dataset in datasets:
             self.append(dataset)
 
-    def to_json(self, file_path: Union[str, pathlib.Path]) -> None:
+    def to_json(self, file_path: str | pathlib.Path) -> None:
         r"""Save the ContinualAlignmentDataset to a json file.
 
         Note: Uses to_dict() under the hood to get a dictionary representation.
@@ -90,8 +87,6 @@ def to_json(self, file_path: Union[str, pathlib.Path]) -> None:
     def to_dict(self) -> Dict[str, Any]:
         r"""Convert the ContinualAlignmentDataset to dictionary represenetation.
 
-        Note: This method is the functional inverse of ContinualAlignmentDataset.from_dict().
-
         Returns:
             Dict[str, Any]: The dictionary representation of the ContinualAlignmentDataset.
         """
@@ -101,9 +96,7 @@ def to_dict(self) -> Dict[str, Any]:
         return dataset_dict
 
     @classmethod
-    def from_json(
-        cls, file_path: Union[str, pathlib.Path]
-    ) -> 'ContinualAlignmentDataset':
+    def from_json(cls, file_path: str | pathlib.Path) -> ContinualAlignmentDataset:
         r"""Load the ContinualAlignmentDataset from a json file.
 
         Note: Uses ContinualAlignmentDataset.from_dict() under the hood to parse the representation.
@@ -119,7 +112,7 @@ def from_json(
         return cls.from_dict(dataset_dict)
 
     @classmethod
-    def from_dict(cls, dataset_dict: Dict[str, Any]) -> 'ContinualAlignmentDataset':
+    def from_dict(cls, dataset_dict: Dict[str, Any]) -> ContinualAlignmentDataset:
         r"""Construct a ContinualAlignmentDataset from dictionary representation.
 
         Note:

aif_gen/dataset/split/functional.py

@@ -23,6 +23,6 @@ def split(
             )
     # just change the fractions
     for i in range(len(dataset.datasets)):
-        dataset.datasets[i]._train_frac = 1 - test_ratio
+        dataset.datasets[i].train_frac = 1 - test_ratio
 
     return dataset

test/test_validation/test_count_validation.py

@@ -4,6 +4,8 @@
     ContinualAlignmentDataset,
 )
 from aif_gen.dataset.validation import count_validation
+from aif_gen.task.alignment_task import AlignmentTask
+from aif_gen.task.domain import Domain
 
 
 def test_count_validation_all_unique():
@@ -18,7 +20,9 @@ def test_count_validation_all_unique():
             'Mock prompt C 1', 'Winning Response C 1', 'Losing Response C 1'
         ),
     ]
-    mock_task = None
+    mock_task = AlignmentTask(
+        domain=Domain.from_dict({'education': {}}), objective='', preference=''
+    )
     dataset = AlignmentDataset(task=mock_task, samples=samples)
     expected_counts = [
         {
@@ -44,7 +48,9 @@ def test_count_validation_all_same_prompts():
             'Mock prompt A 2', 'Winning Response C 2', 'Losing Response C 2'
         ),
     ]
-    mock_task = None
+    mock_task = AlignmentTask(
+        domain=Domain.from_dict({'education': {}}), objective='', preference=''
+    )
     dataset = AlignmentDataset(task=mock_task, samples=samples)
     expected_counts = [
         {
@@ -70,7 +76,9 @@ def test_count_validation_all_same_responses():
             'Mock prompt C 3', 'Winning Response A 3', 'Losing Response B 3'
         ),
     ]
-    mock_task = None
+    mock_task = AlignmentTask(
+        domain=Domain.from_dict({'education': {}}), objective='', preference=''
+    )
     dataset = AlignmentDataset(task=mock_task, samples=samples)
     expected_counts = [
         {
@@ -96,7 +104,9 @@ def test_count_validation_all_same_everything():
             'Mock prompt A 4', 'Winning Response A 4', 'Losing Response A 4'
         ),
     ]
-    mock_task = None
+    mock_task = AlignmentTask(
+        domain=Domain.from_dict({'education': {}}), objective='', preference=''
+    )
     dataset = AlignmentDataset(task=mock_task, samples=samples)
     expected_counts = [
         {
@@ -122,7 +132,9 @@ def test_count_countinual_dataset():
             'Mock prompt C 1', 'Winning Response C 1', 'Losing Response C 1'
         ),
     ]
-    mock_task = None
+    mock_task = AlignmentTask(
+        domain=Domain.from_dict({'education': {}}), objective='', preference=''
+    )
     dataset_one = AlignmentDataset(task=mock_task, samples=samples)
 
     samples = [
@@ -136,7 +148,6 @@ def test_count_countinual_dataset():
             'Mock prompt A 2', 'Winning Response C 2', 'Losing Response C 2'
         ),
     ]
-    mock_task = None
     dataset_two = AlignmentDataset(task=mock_task, samples=samples)
 
     samples = [
@@ -150,7 +161,6 @@ def test_count_countinual_dataset():
             'Mock prompt C 3', 'Winning Response A 3', 'Losing Response B 3'
         ),
     ]
-    mock_task = None
     dataset_three = AlignmentDataset(task=mock_task, samples=samples)
 
     samples = [
@@ -164,7 +174,6 @@ def test_count_countinual_dataset():
             'Mock prompt A 4', 'Winning Response A 4', 'Losing Response A 4'
         ),
     ]
-    mock_task = None
     dataset_four = AlignmentDataset(task=mock_task, samples=samples)
 
     dataset = ContinualAlignmentDataset(
@@ -216,7 +225,9 @@ def test_count_validation_stop_words_removed():
             'with Mock prompt A 4', 'by Winning Response A 4', 'is Losing Response A 4'
         ),
     ]
-    mock_task = None
+    mock_task = AlignmentTask(
+        domain=Domain.from_dict({'education': {}}), objective='', preference=''
+    )
     dataset = AlignmentDataset(task=mock_task, samples=samples)
     expected_counts = [
         {