Merge pull request #22 from ComplexData-MILA/dev/alignment_dataset

Alignment Dataset Object

Author: Jacob-Chmura

aif_gen/dataset/__init__.py

@@ -0,0 +1,3 @@
+from aif_gen.dataset.alignment_sample import AlignmentDatasetSample
+from aif_gen.dataset.alignment_dataset import AlignmentDataset
+from aif_gen.dataset.continual_alignment_dataset import ContinualAlignmentDataset

aif_gen/dataset/alignment_dataset.py

@@ -0,0 +1,145 @@
+import json
+from dataclasses import asdict
+from typing import Any, Dict, List, Union
+
+from aif_gen.task import AlignmentTask
+
+from .alignment_sample import AlignmentDatasetSample
+
+
+class AlignmentDataset:
+    r"""Container object for an Alignment Dataset.
+
+    Args:
+        task (AligmnentTask): The AlignmentTask associated with the dataset.
+        samples (List[AlignmentDatasetSample]): The samples in this AlignmentDataset.
+    """
+
+    def __init__(
+        self, task: AlignmentTask, samples: List[AlignmentDatasetSample]
+    ) -> None:
+        self._task = task
+        self._samples = samples
+
+    @property
+    def task(self) -> AlignmentTask:
+        """AlignmentTask: The task associated with the AlignmentDataset."""
+        return self._task
+
+    @property
+    def samples(self) -> List[AlignmentDatasetSample]:
+        """List[AlignmentDatasetSample]: The list of samples associated with the AlignmentDataset."""
+        return self._samples
+
+    @property
+    def num_samples(self) -> int:
+        """int: The number of samples associated with the AlignmentDataset."""
+        return len(self.samples)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(
+        self, key: Union[slice, int]
+    ) -> Union[AlignmentDatasetSample, List[AlignmentDatasetSample]]:
+        # Slicing directly on the samples
+        return self.samples[key]
+
+    def append(self, sample: AlignmentDatasetSample) -> None:
+        r"""Append a single AlignmentDatasetSample to the Alignment Dataset.
+
+        Args:
+            sample (AlignmentDatasetSample): The new sample to add.
+
+        Raises:
+            TypeError: if the sample is not of type AlignmentDatasetSample.
+        """
+        if isinstance(sample, AlignmentDatasetSample):
+            self.samples.append(sample)
+        else:
+            raise TypeError(
+                f'Sample: {sample} must be of type AlignmentDatasetSample but got {sample.__class__.__name__}'
+            )
+
+    def extend(self, samples: List[AlignmentDatasetSample]) -> None:
+        r"""Add multiple AlignmentDatasetSample's to the Alignment Dataset.
+
+        Args:
+            samples (List[AlignmentDatasetSample]): The new samples to add.
+
+        Raises:
+            TypeError: if any sample is not of type AlignmentDatasetSample.
+        """
+        for sample in samples:
+            self.append(sample)
+
+    def to_json(self, file_path: str) -> None:
+        r"""Save the AlignmentDataset to a json file.
+
+        Note: Uses to_dict() under the hood to get a dictionary representation.
+
+        Args:
+            file_path (str): The os.pathlike object to write to.
+        """
+        dataset_dict = self.to_dict()
+        with open(file_path, 'w') as f:
+            json.dump(dataset_dict, f)
+
+    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['samples'] = []
+        for sample in self.samples:
+            dataset_dict['samples'].append(asdict(sample))
+        return dataset_dict
+
+    @classmethod
+    def from_json(cls, file_path: str) -> 'AlignmentDataset':
+        r"""Load the AlignmentDataset from a json file.
+
+        Note: Uses AlignmentDataset.from_dict() under the hood to parse the representation.
+
+        Args:
+            file_path (str): The os.pathlike object to read from.
+
+        Returns:
+            AlignmentDataset: The newly constructed 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':
+        r"""Construct an AlignmentDataset from dictionary representation.
+
+        Note:
+            Expects 'task', and 'samples' keys to be present in the dictionary.
+            The 'task' value should be parsable by AlignmentTask.from_dict().
+            The 'samples' value should be a list of dictionaries, each of which
+            are parsable by AlignmentDatasetSample.
+
+        Args:
+            dataset_dict (Dict[str, Any]): The dictionary that encodes the AlignmentDataset.
+
+        Returns:
+            AlignmentDataset: The newly constructed AlignmentDataset.
+
+        Raises:
+            ValueError: If the input dictionary is missing any required keys.
+        """
+        task = AlignmentTask.from_dict(dataset_dict['task'])
+        samples = []
+        for sample in dataset_dict['samples']:
+            sample = AlignmentDatasetSample(**sample)
+            samples.append(sample)
+
+        return cls(task, samples)

aif_gen/dataset/alignment_sample.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class AlignmentDatasetSample:
+    r"""Container for a single Alignment Dataset Sample.
+
+    This representation is faithful to the "TRL Preference Format with explicit prompt".
+    See: https://huggingface.co/docs/trl/en/dataset_formats.
+
+    Args:
+        prompt (str): The prompt associated with the sample.
+        chosen (str): The winning response associated with the sample.
+        rejected (str): The losing response associated with the sample.
+    """
+
+    prompt: str
+    chosen: str
+    rejected: str

aif_gen/dataset/continual_alignment_dataset.py

@@ -0,0 +1,136 @@
+import json
+from typing import Any, Dict, List, Union
+
+from aif_gen.dataset.alignment_sample import AlignmentDatasetSample
+
+from .alignment_dataset import AlignmentDataset
+
+
+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]:
+        """List[AlignmentDataset]: The list of AlignmentDataset constituents."""
+        return self._datasets
+
+    @property
+    def num_datasets(self) -> int:
+        """int: The number of AlignmentDataset constituents."""
+        return len(self.datasets)
+
+    @property
+    def num_samples(self) -> int:
+        """int: The total number of samples acros all AlignmentDataset constituents."""
+        return sum(len(dataset) for dataset in self.datasets)
+
+    def __len__(self) -> int:
+        """int: The total number of samples acros all AlignmentDataset constituents."""
+        return self.num_samples
+
+    def __getitem__(
+        self, key: Union[slice, int]
+    ) -> Union[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.
+
+        Args:
+            dataset (AlignmentDataset): The new dataset to add.
+
+        Raises:
+            TypeError: if the sample is not of type AlignmentDataset.
+        """
+        if isinstance(dataset, AlignmentDataset):
+            self.datasets.append(dataset)
+        else:
+            raise TypeError(
+                f'Dataset: {dataset} must be of type AlignmentDataset but got {dataset.__class__.__name__}'
+            )
+
+    def extend(self, datasets: List[AlignmentDataset]) -> None:
+        r"""Append multiple AlignmentDataset's to the ConitnualAlignmentDataset.
+
+        Args:
+            datasets (List[AlignmentDataset]): The new datasets to add.
+
+        Raises:
+            TypeError: if any dataset is not of type AlignmentDataset.
+        """
+        for dataset in datasets:
+            self.append(dataset)
+
+    def to_json(self, file_path: str) -> None:
+        r"""Save the ContinualAlignmentDataset to a json file.
+
+        Note: Uses to_dict() under the hood to get a dictionary representation.
+
+        Args:
+            file_path (str): The os.pathlike object to write to.
+        """
+        dataset_dict = self.to_dict()
+        with open(file_path, 'w') as f:
+            json.dump(dataset_dict, f)
+
+    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.
+        """
+        dataset_dict: Dict[str, List[Any]] = {'datasets': []}
+        for dataset in self.datasets:
+            dataset_dict['datasets'].append(dataset.to_dict())
+        return dataset_dict
+
+    @classmethod
+    def from_json(cls, file_path: str) -> 'ContinualAlignmentDataset':
+        r"""Load the ContinualAlignmentDataset from a json file.
+
+        Note: Uses ContinualAlignmentDataset.from_dict() under the hood to parse the representation.
+
+        Args:
+            file_path (str): The os.pathlike object to read from.
+
+        Returns:
+            ContinualAlignmentDataset: The newly constructed ContinualAlignmentDataset.
+        """
+        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]) -> 'ContinualAlignmentDataset':
+        r"""Construct a ContinualAlignmentDataset from dictionary representation.
+
+        Note:
+            Expects 'datasets' key to be present in the dictionary. The value is a list
+            of dictionaries, each parsable by AlignmentDataset.from_dict().
+
+        Args:
+            dataset_dict (Dict[str, Any]): The dictionary that encodes the ContinualAlignmentDataset.
+
+        Returns:
+            ContinualAlignmentDataset: The newly constructed ContinualAlignmentDataset.
+
+        Raises:
+            ValueError: If the input dictionary is missing any required keys.
+        """
+        datasets = []
+        for dataset in dataset_dict['datasets']:
+            datasets.append(AlignmentDataset.from_dict(dataset))
+        return cls(datasets)