camel-ai 0.2.28__py3-none-any.whl → 0.2.29__py3-none-any.whl

camel/datasets/base.py

@@ -13,28 +13,26 @@
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
 
 import json
-import os
 import random
+from datetime import datetime
 from pathlib import Path
 from typing import (
     Any,
-    Callable,
     Dict,
     List,
     Optional,
     Sized,
-    TypeVar,
     Union,
 )
 
-import torch
 from datasets import Dataset as HFDataset
 from pydantic import BaseModel, Field, ValidationError
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import Dataset
 
 from camel.agents import ChatAgent
 from camel.logger import get_logger
 from camel.verifiers import BaseVerifier
+from camel.verifiers.models import VerifierInput
 
 logger = get_logger(__name__)
 
@@ -44,13 +42,9 @@ class DataPoint(BaseModel):
 
     Attributes:
         question (str): The primary question or issue to be addressed.
-        rationale (str): Logical reasoning or explanation behind the
-            answer.
+        rationale (Optional[str]): Logical reasoning or explanation behind the
+            answer. (default: :obj:`None`)
         final_answer (str): The final answer.
-        raw_markdown (Optional[str]): Raw markdown content for generating
-            rewards/hints. (default: :obj:`None`)
-        difficulty (Optional[str]): Difficulty level of the question.
-            (default: :obj:`None`)
         metadata Optional[Dict[str, Any]]: Additional metadata about the data
             point. (default: :obj:`None`)
     """
@@ -58,13 +52,12 @@ class DataPoint(BaseModel):
     question: str = Field(
         ..., description="The primary question or issue to be addressed."
     )
-    rationale: str = Field(
-        ..., description="Logical reasoning or explanation behind the answer."
+    rationale: Optional[str] = Field(
+        default=None,
+        description="Logical reasoning or explanation behind the answer.",
     )
     final_answer: str = Field(..., description="The final answer.")
-    difficulty: Optional[str] = Field(
-        None, description="Difficulty level of the question."
-    )
+
     metadata: Optional[Dict[str, Any]] = Field(
         default=None, description="Additional metadata about the data point."
     )
@@ -90,244 +83,11 @@ class DataPoint(BaseModel):
         return cls(**data)
 
 
-class BaseDataset(Dataset):
-    r"""A dataset contains questions and ground truth data for training.
-    It can be either static (e.g., MATH dataset) or generative
-    (using an LLM to generate questions).
-    """
-
-    def __init__(
-        self,
-        data: List[Dict[str, str]],
-        cache_dir: Optional[str] = None,
-        **kwargs,
-    ):
-        r"""Initialize the dataset.
-
-        Args:
-            data (List[Dict[str, str]]): List of dictionary items to
-                create the dataset from.
-            cache_dir (Optional[str]): Directory to cache dataset files.
-                (default: :obj:`None`)
-            **kwargs: Additional dataset parameters.
-
-        Note:
-            The dataset must be initialized by calling setup() before use.
-        """
-        self._is_setup = False
-        self._raw_data: List[Dict[str, str]] = data if data is not None else []
-        self._cache_dir = str(cache_dir) if cache_dir is not None else None
-
-        # Store all parameters in metadata dict for compatibility
-        self._metadata = {
-            'cache_dir': self._cache_dir,
-            **kwargs,
-        }
-
-        self.data: List[DataPoint] = []  # Will be populated in setup()
-
-    async def setup(self) -> None:
-        r"""Set up the dataset with necessary resources.
-
-        This method:
-        1. Creates cache directory if needed
-        2. Processes raw data into DataPoint objects using vectorized
-        operations
-        3. Validates dataset integrity
-
-        Raises:
-            OSError: If cache directory creation fails.
-            Exception: If dataset initialization fails.
-        """
-        if self._is_setup:
-            logger.debug(f"{self.__class__.__name__} already initialized")
-            return
-
-        try:
-            # Create cache directory if specified
-            if self._cache_dir:
-                try:
-                    os.makedirs(self._cache_dir, exist_ok=True)
-                    logger.debug(f"Created cache directory: {self._cache_dir}")
-                except OSError as e:
-                    logger.error(
-                        f"Failed to create cache directory "
-                        f"{self._cache_dir}: {e}"
-                    )
-                    raise
-
-            # Process raw data into DataPoint objects using vectorized
-            # operations
-            if not self._raw_data:
-                self.data = []
-                logger.debug("No raw data to process")
-            else:
-                try:
-                    # Helper function for validation that can be used with map
-                    def create_datapoint(item, idx=None):
-                        try:
-                            return DataPoint(
-                                question=item.get('question', ''),
-                                rationale=item.get('rationale', ''),
-                                final_answer=item.get('final_answer', ''),
-                                metadata=item.get('metadata', {})
-                                if isinstance(item.get('metadata'), dict)
-                                else {},
-                                raw_markdown='',
-                                difficulty='',
-                            )
-                        except ValidationError as e:
-                            idx_str = (
-                                f" at index {idx}" if idx is not None else ""
-                            )
-                            error_msg = (
-                                f"Sample{idx_str} validation error: {e}"
-                            )
-                            logger.error(error_msg)
-                            raise ValueError(error_msg)
-
-                    # If raw_data is already a HF dataset, use its map function
-                    if hasattr(self._raw_data, 'map') and callable(
-                        self._raw_data.map
-                    ):
-                        # Using HF dataset's map for vectorized processing
-                        processed_data = self._raw_data.map(
-                            lambda example, idx: {
-                                'datapoint': create_datapoint(example, idx)
-                            },
-                            with_indices=True,
-                        )
-                        self.data = [
-                            item['datapoint'] for item in processed_data
-                        ]
-                    else:
-                        # Bulk create datapoints
-                        self.data = [
-                            create_datapoint(item, i)
-                            for i, item in enumerate(self._raw_data)
-                        ]
-
-                    logger.debug(f"Processed {len(self.data)} data points")
-                except Exception as e:
-                    logger.error(f"Error processing data: {e}")
-                    raise
-
-            self._is_setup = True
-            logger.info(f"{self.__class__.__name__} initialized successfully")
-
-        except Exception as e:
-            logger.error(f"Error during {self.__class__.__name__} setup: {e}")
-            await self.cleanup()
-            raise
-
-    async def cleanup(self) -> None:
-        r"""Clean up dataset resources.
-
-        This method handles cleanup of resources and resets the dataset state.
-        """
-        if not self._is_setup:
-            return
-
-        try:
-            # Clear metadata while preserving init config
-            init_config = {
-                'cache_dir': self._cache_dir,
-            }
-            self._metadata = init_config
-
-            logger.info(f"{self.__class__.__name__} cleaned up successfully")
-
-        except Exception as e:
-            logger.error(
-                f"Error during {self.__class__.__name__} cleanup: {e}"
-            )
-            raise
-
-        finally:
-            # Always mark as uninitialized, even if cleanup fails
-            self._is_setup = False
-
-    def sample(self) -> DataPoint:
-        r"""Sample a random datapoint from the dataset.
-
-        Returns:
-            DataPoint: A randomly sampled DataPoint.
-
-        Raises:
-            RuntimeError: If dataset is not initialized.
-        """
-        if not self._is_setup:
-            raise RuntimeError(
-                f"{self.__class__.__name__} must be initialized "
-                "before sampling"
-            )
-
-        idx = random.randint(0, len(self) - 1)
-        return self[idx]
-
-    def __len__(self) -> int:
-        r"""Return the size of the dataset."""
-        return len(self.data)
-
-    def __getitem__(self, idx: int) -> DataPoint:
-        r"""Get an item from the dataset.
-
-        Args:
-            idx (int): Index of the item to get.
-
-        Returns:
-            DataPoint: DataPoint from the dataset with the given index.
-
-        Raises:
-            IndexError: If idx is out of bounds.
-        """
-        if idx < 0 or idx >= len(self):
-            raise IndexError(
-                f"Index {idx} out of bounds for dataset of size {len(self)}"
-            )
-
-        return self.data[idx]
-
-    @property
-    def metadata(self) -> Dict[str, Any]:
-        r"""Get dataset metadata."""
-        return self._metadata.copy()
-
-    def to_pytorch_dataset(
-        self,
-        transform: Optional[Callable] = None,
-        target_transform: Optional[Callable] = None,
-        batch_size: Optional[int] = None,
-    ) -> Union["PyTorchDataset", "DataLoader"]:
-        r"""Convert to a PyTorch dataset or DataLoader.
-
-        Args:
-            transform (Optional[Callable]): Transform to apply to samples.
-            target_transform (Optional[Callable]): Transform to apply to
-                targets.
-            batch_size (Optional[int]): If provided, returns a DataLoader with
-                the specified batch size instead of a PyTorchDataset.
-                (default: :obj:`None`)
-
-        Returns:
-            Union[PyTorchDataset, torch.utils.data.DataLoader]: Dataset in
-                PyTorch format or DataLoader if batch_size is provided.
-        """
-        dataset = PyTorchDataset.from_datapoints(
-            self.data,
-            transform=transform,
-            target_transform=target_transform,
-        )
-
-        if batch_size is not None:
-            return dataset.get_dataloader(batch_size=batch_size)
-
-        return dataset
-
-
-class SeedDataset(BaseDataset):
-    r"""A dataset containing validated seed examples for data generation.
+class StaticDataset(Dataset):
+    r"""A static dataset containing a list of datapoints.
     Ensures that all items adhere to the DataPoint schema.
+    This dataset extends :obj:`Dataset` from PyTorch and should
+    be used when its size is fixed at runtime.
 
     This class can initialize from Hugging Face Datasets,
     PyTorch Datasets, JSON file paths, or lists of dictionaries,
@@ -337,46 +97,46 @@ class SeedDataset(BaseDataset):
     def __init__(
         self,
         data: Union[HFDataset, Dataset, Path, List[Dict[str, Any]]],
-        cache_dir: Optional[str] = None,
-        seed: Optional[int] = None,
+        seed: int = 42,
         min_samples: int = 1,
         strict: bool = False,
         **kwargs,
     ):
-        r"""Initialize the seed dataset and validate integrity.
+        r"""Initialize the static dataset and validate integrity.
 
         Args:
-            data (Union[HFDataset, Dataset, Path, List[Dict[str, Any]]]):
-            Input data, which can be:
-                - A Hugging Face Dataset (HFDataset)
-                - A PyTorch Dataset (torch.utils.data.Dataset)
-                - A Path object representing the path to a JSON file
-                - A list of dictionaries with DataPoint-compatible fields
-            seed (Optional[int]): Seed for reproducibility.
-                (default: :obj:`1`)
-            min_samples (int): Minimum number of samples required.
-                (default: :obj:`1`)
-            strict (bool): Whether to raise an error on invalid datapoints
-                (True) or skip/filter them (False). (default: False)
+            data (:obj:`Union[HFDataset, Dataset,
+            Path, List[Dict[str, Any]]]`):
+                Input data, which can be one of the following:
+                - A Hugging Face Dataset (:obj:`HFDataset`).
+                - A PyTorch Dataset (:obj:`torch.utils.data.Dataset`).
+                - A :obj:`Path` object representing a JSON file.
+                - A list of dictionaries with :obj:`DataPoint`-compatible
+                  fields.
+            seed (:obj:`int`): Random seed for reproducibility.
+                Default is :obj:`42`.
+            min_samples (:obj:`int`): Minimum required number of samples.
+                Default is :obj:`1`.
+            strict (:obj:`bool`): Whether to raise an error on invalid
+                datapoints (:obj:`True`) or skip/filter them (:obj:`False`).
+                Default is :obj:`False`.
             **kwargs: Additional dataset parameters.
 
         Raises:
-            TypeError: If the data type is not supported.
-            ValueError: If dataset size is less than min_samples or
-            if sample validation fails.
-            FileNotFoundError: If the JSON file path doesn't exist.
-            json.JSONDecodeError: If the JSON file is invalid.
+            TypeError: If the input data type is unsupported.
+            ValueError: If the dataset contains fewer than :obj:`min_samples`
+                datapoints or if validation fails.
+            FileNotFoundError: If the specified JSON file path does not exist.
+            json.JSONDecodeError: If the JSON file contains invalid formatting.
         """
-        # Initialize BaseDataset with empty data, we'll populate it ourselves
-        super().__init__(data=[], cache_dir=cache_dir, **kwargs)
 
+        # Store all parameters in metadata dict for compatibility
+        self._metadata = {
+            **kwargs,
+        }
         self._rng = random.Random(seed)
         self._strict = strict
 
-        # Type checking and conversion into list of dicts to have a
-        # consistent internal format. Since Seed Dataset should be
-        # small, we can load it entirely into memory
-
         self.data: List[DataPoint] = self._init_data(data)
         self._length = len(self.data)
 
@@ -389,6 +149,26 @@ class SeedDataset(BaseDataset):
     def _init_data(
         self, data: Union[HFDataset, Dataset, Path, List[Dict[str, Any]]]
     ) -> List[DataPoint]:
+        r"""Convert input data from various formats into a list of
+        :obj:`DataPoint` instances.
+
+        Args:
+            data (:obj:`Union[
+                HFDataset,
+                Dataset,
+                Path,
+                List[Dict[str, Any]]
+            ]`):
+                Input dataset in one of the supported formats.
+
+        Returns:
+            :obj:`List[DataPoint]`: A list of validated :obj:`DataPoint`
+                instances.
+
+        Raises:
+            TypeError: If the input data type is unsupported.
+        """
+
         if isinstance(data, HFDataset):
             raw_data = self._init_from_hf_dataset(data)
         elif isinstance(data, Dataset):
@@ -450,7 +230,6 @@ class SeedDataset(BaseDataset):
                     rationale=rationale,
                     final_answer=final_answer,
                     metadata=item.get('metadata'),
-                    difficulty=item.get('difficulty'),
                 )
             except ValidationError as e:
                 if self._strict:
@@ -474,17 +253,19 @@ class SeedDataset(BaseDataset):
         return self._length
 
     def __getitem__(self, idx: int) -> DataPoint:
-        r"""Get an item from the dataset.
+        r"""Retrieve a datapoint by index.
 
         Args:
-            idx (int): Index of the item to get.
+            idx (:obj:`int`): Index of the datapoint.
 
         Returns:
-            DataPoint: DataPoint from the dataset with the given index.
+            :obj:`DataPoint`: The datapoint corresponding to the given index.
 
         Raises:
-            IndexError: If idx is out of bounds.
+            IndexError: If :obj:`idx` is out of bounds (negative or greater
+                than dataset length - 1).
         """
+
         if idx < 0 or idx >= self._length:
             raise IndexError(
                 f"Index {idx} out of bounds for dataset of size {self._length}"
@@ -495,11 +276,12 @@ class SeedDataset(BaseDataset):
         r"""Sample a random datapoint from the dataset.
 
         Returns:
-            DataPoint: A randomly sampled DataPoint.
+            :obj:`DataPoint`: A randomly sampled :obj:`DataPoint`.
 
         Raises:
-            RuntimeError: If the dataset is empty.
+            RuntimeError: If the dataset is empty and no samples can be drawn.
         """
+
         if self._length == 0:
             raise RuntimeError("Dataset is empty, cannot sample.")
         idx = self._rng.randint(0, self._length - 1)
@@ -507,15 +289,42 @@ class SeedDataset(BaseDataset):
 
     @property
     def metadata(self) -> Dict[str, Any]:
-        r"""Get dataset metadata."""
+        r"""Retrieve dataset metadata.
+
+        Returns:
+            :obj:`Dict[str, Any]`: A copy of the dataset metadata dictionary.
+        """
+
         return self._metadata.copy()
 
     def _init_from_hf_dataset(self, data: HFDataset) -> List[Dict[str, Any]]:
+        r"""Convert a Hugging Face dataset into a list of dictionaries.
+
+        Args:
+            data (:obj:`HFDataset`): A Hugging Face dataset.
+
+        Returns:
+            :obj:`List[Dict[str, Any]]`: A list of dictionaries representing
+            the dataset, where each dictionary corresponds to a datapoint.
+        """
         return [dict(item) for item in data]
 
     def _init_from_pytorch_dataset(
         self, data: Dataset
     ) -> List[Dict[str, Any]]:
+        r"""Convert a PyTorch dataset into a list of dictionaries.
+
+        Args:
+            data (:obj:`Dataset`): A PyTorch dataset.
+
+        Returns:
+            :obj:`List[Dict[str, Any]]`: A list of dictionaries representing
+            the dataset.
+
+        Raises:
+            TypeError: If the dataset does not implement :obj:`__len__()`
+                or contains non-dictionary elements.
+        """
         if not isinstance(data, Sized):
             raise TypeError(
                 f"{type(data).__name__} does not implement `__len__()`."
@@ -533,6 +342,20 @@ class SeedDataset(BaseDataset):
         return raw_data
 
     def _init_from_json_path(self, data: Path) -> List[Dict[str, Any]]:
+        r"""Load and parse a dataset from a JSON file.
+
+        Args:
+            data (:obj:`Path`): Path to the JSON file.
+
+        Returns:
+            :obj:`List[Dict[str, Any]]`: A list of datapoint dictionaries.
+
+        Raises:
+            FileNotFoundError: If the specified JSON file does not exist.
+            ValueError: If the JSON content is not a list of dictionaries.
+            json.JSONDecodeError: If the JSON file has invalid formatting.
+        """
+
         if not data.exists():
             raise FileNotFoundError(f"JSON file not found: {data}")
         try:
@@ -557,6 +380,18 @@ class SeedDataset(BaseDataset):
     def _init_from_list(
         self, data: List[Dict[str, Any]]
     ) -> List[Dict[str, Any]]:
+        r"""Validate and convert a list of dictionaries into a dataset.
+
+        Args:
+            data (:obj:`List[Dict[str, Any]]`): A list of dictionaries where
+                each dictionary must be a valid :obj:`DataPoint`.
+
+        Returns:
+            :obj:`List[Dict[str, Any]]`: The validated list of dictionaries.
+
+        Raises:
+            ValueError: If any item in the list is not a dictionary.
+        """
         for i, item in enumerate(data):
             if not isinstance(item, dict):
                 raise ValueError(
@@ -566,111 +401,7 @@ class SeedDataset(BaseDataset):
         return data
 
 
-class SyntheticDataset(BaseDataset):
-    r"""A dataset for storing synthetically generated data points.
-
-    This class is used to store datapoints that are generated through
-    a generative process, such as using an agent.
-    """
-
-    def __init__(
-        self,
-        data: Optional[List[Dict[str, str]]] = None,
-        cache_dir: Optional[str] = None,
-        **kwargs,
-    ):
-        r"""Initialize the synthetic dataset.
-
-        Args:
-            data (Optional[List[Dict[str, str]]]): List of dictionary items to
-                create the dataset from. (default: :obj:`None`)
-            cache_dir (Optional[str]): Directory to cache dataset files.
-                (default: :obj:`None`)
-            **kwargs: Additional dataset parameters.
-        """
-        super().__init__(
-            data=data if data is not None else [],
-            cache_dir=cache_dir,
-            **kwargs,
-        )
-        self.data: List[DataPoint] = []
-
-    def add(self, item: DataPoint) -> None:
-        r"""Add a new data point to the dataset.
-
-        Args:
-            item (DataPoint): The datapoint to add to the dataset.
-        """
-        self.data.append(item)
-
-    def add_batch(self, items: List[DataPoint]) -> None:
-        r"""Add multiple data points to the dataset.
-
-        Args:
-            items (List[DataPoint]): The datapoints to add to the dataset.
-        """
-        self.data.extend(items)
-
-    def to_pytorch_dataset(
-        self,
-        transform: Optional[Callable] = None,
-        target_transform: Optional[Callable] = None,
-        batch_size: Optional[int] = None,
-    ) -> Union["PyTorchDataset", "DataLoader"]:
-        r"""Convert to a PyTorch dataset or DataLoader.
-
-        Args:
-            transform (Optional[Callable]): Transform to apply to samples.
-            target_transform (Optional[Callable]): Transform to apply to
-                targets.
-            batch_size (Optional[int]): If provided, returns a DataLoader with
-                the specified batch size instead of a PyTorchDataset.
-                (default: :obj:`None`)
-
-        Returns:
-            Union[PyTorchDataset, torch.utils.data.DataLoader]: Dataset in
-                PyTorch format or DataLoader if batch_size is provided.
-        """
-        return convert_synthetic_to_pytorch(
-            self,
-            transform=transform,
-            target_transform=target_transform,
-            batch_size=batch_size,
-        )
-
-    def save_pytorch_format(self, path: str, compression: bool = True) -> None:
-        r"""Save the dataset to disk in PyTorch format.
-
-        Args:
-            path (str): Path to save the dataset to.
-            compression (bool): Whether to use compression to reduce file size.
-                (default: :obj:`True`)
-        """
-        save_synthetic_dataset(self, path, compression=compression)
-
-    def filter(
-        self, predicate: Callable[[DataPoint], bool]
-    ) -> 'SyntheticDataset':
-        r"""Filter the dataset using a predicate function.
-
-        Args:
-            predicate (Callable[[DataPoint], bool]): Function that takes a
-                DataPoint and returns True if it should be kept, False
-                otherwise.
-
-        Returns:
-            SyntheticDataset: A new dataset containing only the filtered items.
-        """
-        filtered_data = [dp for dp in self.data if predicate(dp)]
-
-        # Create a new dataset with the filtered data
-        new_dataset = SyntheticDataset()
-        new_dataset.add_batch(filtered_data)
-
-        return new_dataset
-
-
-class GenerativeDataset(BaseDataset):
+class GenerativeDataset(Dataset):
     r"""A dataset for generating synthetic datapoints using external agents and
     verifiers.
 
@@ -680,30 +411,22 @@ class GenerativeDataset(BaseDataset):
 
     def __init__(
         self,
-        seed_dataset: SeedDataset,
+        seed_dataset: StaticDataset,
         verifier: BaseVerifier,
         agent: ChatAgent,
-        cache_dir: Optional[str] = None,
         seed: int = 42,
         **kwargs,
     ):
         r"""Initialize the generative dataset.
 
         Args:
-            seed_dataset (SeedDataset): Validated dataset to use for examples.
+            seed_dataset (StaticDataset): Validated static dataset to
+            use for examples.
             verifier (BaseVerifier): Verifier to validate generated content.
             agent (ChatAgent): Agent to generate new datapoints.
-            cache_dir (Optional[str]): Directory to cache dataset files.
-                (default: :obj:`None`)
             seed (int): Random seed for reproducibility. (default: :obj:`42`)
             **kwargs: Additional dataset parameters.
         """
-        # Initialize with empty data since we'll generate content dynamically
-        super().__init__(
-            data=[],
-            cache_dir=cache_dir,
-            **kwargs,
-        )
 
         self.seed_dataset = seed_dataset
         self.verifier = verifier
@@ -712,8 +435,11 @@ class GenerativeDataset(BaseDataset):
         self.seed = seed
         random.seed(self.seed)
 
+        self._data: List[DataPoint] = []
+
     def _construct_prompt(self, examples: List[DataPoint]) -> str:
-        r"""Construct a prompt for generating new datapoints.
+        r"""Construct a prompt for generating new datapoints
+        using a fixed sample of 3 examples from the seed dataset.
 
         Args:
             examples (List[DataPoint]): Examples to include in the prompt.
@@ -732,440 +458,182 @@ class GenerativeDataset(BaseDataset):
         prompt += "New datapoint:"
         return prompt
 
-    async def generate_new(self, n: int) -> None:
-        r"""Generate n new datapoints and add them to the dataset.
+    async def generate_new(
+        self, n: int, max_retries: int = 10
+    ) -> List[DataPoint]:
+        r"""Generates and validates `n` new datapoints through
+        few-shot prompting, with a retry limit.
+
+        Steps:
+            1. Samples 3 examples from the seed dataset.
+            2. Constructs a prompt using the selected examples.
+            3. Uses an agent to generate a new datapoint.
+            4. Verifies the datapoint using a verifier.
+            5. Stores valid datapoints in memory.
 
         Args:
             n (int): Number of valid datapoints to generate.
+            max_retries (int): Maximum number of retries before stopping.
+
+        Returns:
+            List[DataPoint]: A list of newly generated valid datapoints.
 
-        This method generates new datapoints by:
-        1. Sampling examples from the seed dataset
-        2. Constructing a prompt for the agent
-        3. Generating a new datapoint using the agent
-        4. Verifying the generated datapoint with the verifier
-        5. Adding valid datapoints to the dataset
+        Raises:
+            TypeError: If the agent's output is not a dictionary (or does not
+            match the expected format).
+            KeyError: If required keys are missing from the response.
+            AttributeError: If the verifier response lacks attributes.
+            ValidationError: If a datapoint fails schema validation.
+            RuntimeError: If retries are exhausted before `n` valid datapoints
+            are generated.
+
+        Notes:
+            - Retries on validation failures until `n` valid datapoints exist
+            or `max_retries` is reached, whichever comes first.
+            - If retries are exhausted before reaching `n`, a `RuntimeError`
+            is raised.
+            - Metadata includes a timestamp for tracking datapoint creation.
+            - This method can be overridden to implement custom generation.
         """
         valid_data_points: List[DataPoint] = []
+        retries = 0
 
-        while len(valid_data_points) < n:
+        while len(valid_data_points) < n and retries < max_retries:
             try:
-                indices = random.sample(range(len(self.seed_dataset)), 3)
-                examples = [self.seed_dataset[i] for i in indices]
+                examples = [self.seed_dataset.sample() for _ in range(3)]
                 prompt = self._construct_prompt(examples)
 
-                # Get agent response
-                agent_output = (
-                    self.agent.step(prompt, response_format=DataPoint)
-                    .msgs[0]
-                    .parsed
-                )
-
-                if not isinstance(agent_output, dict):
-                    raise TypeError("Agent output must be a dictionary")
-                if (
-                    'question' not in agent_output
-                    or 'rationale' not in agent_output
-                ):
-                    raise KeyError(
-                        "Agent output missing required keys: "
-                        "'question' or 'rationale'"
+                try:
+                    agent_output = (
+                        self.agent.step(prompt, response_format=DataPoint)
+                        .msgs[0]
+                        .parsed
                     )
-
-                rationale = agent_output['rationale']
-
-                # Verify the generated content
-                verifier_response = await self.verifier.verify(rationale)
-                if not hasattr(verifier_response, 'content'):
-                    raise AttributeError(
-                        "Verifier response missing 'content' attribute"
+                    if not isinstance(agent_output, dict):
+                        raise TypeError("Agent output must be a dictionary")
+                    if (
+                        "question" not in agent_output
+                        or "rationale" not in agent_output
+                    ):
+                        raise KeyError(
+                            "Missing 'question' or 'rationale' in agent output"
+                        )
+                except (TypeError, KeyError) as e:
+                    logger.warning(
+                        f"Agent output issue: {e}, retrying... "
+                        f"({retries + 1}/{max_retries})"
                     )
-
-                if not verifier_response.result:
+                    retries += 1
                     continue
 
-                final_answer = verifier_response.result
-
-                # Create and validate the new datapoint
-                new_datapoint = {
-                    'question': agent_output['question'],
-                    'rationale': rationale,
-                    'final_answer': final_answer,
-                }
-
-                datapoint = DataPoint(**new_datapoint)
-                valid_data_points.append(datapoint)
-
-            except (TypeError, KeyError, AttributeError, ValidationError) as e:
-                logger.warning(
-                    f"Error encountered during generation: {e}, retrying..."
-                )
-
-        # Add all valid datapoints to the dataset
-        for datapoint in valid_data_points:
-            self.data.append(datapoint)
-            logger.debug("Added new datapoint to dataset")
-
-
-# Define a type variable for return type flexibility
-T_co = TypeVar('T_co', covariant=True)
-
+                rationale = agent_output["rationale"]
 
-class PyTorchDataset(Dataset[T_co]):
-    r"""A PyTorch-compatible dataset implementation that leverages PyTorch's
-    efficient data handling capabilities.
-    """
-
-    def __init__(
-        self,
-        data: List[Dict[str, Any]],
-        transform: Optional[Callable] = None,
-        target_transform: Optional[Callable] = None,
-        validate: bool = True,
-    ):
-        r"""Initialize the PyTorch dataset.
-
-        Args:
-            data (List[Dict[str, Any]]): List of dictionary items to create
-                the dataset from.
-            transform (Optional[Callable]): A function/transform that takes a
-                sample and returns a transformed version for features.
-                (default: :obj:`None`)
-            target_transform (Optional[Callable]): A function/transform that
-                takes a target and returns a transformed version. (default:
-                :obj:`None`)
-            validate (bool): Whether to validate data points using DataPoint
-                schema. (default: :obj:`True`)
-
-        Raises:
-            ValidationError: If validation is enabled and data doesn't match
-                DataPoint schema.
-        """
-        self.transform = transform
-        self.target_transform = target_transform
-
-        # Validate and store data
-        self._raw_data = data
-        self.data = []
+                try:
+                    verifier_response = await self.verifier.verify(
+                        VerifierInput(
+                            llm_response=rationale, ground_truth=None
+                        )
+                    )
+                    if not verifier_response or not verifier_response.result:
+                        raise ValueError(
+                            "Verifier unsuccessful, response: "
+                            f"{verifier_response}"
+                        )
+                except (ValueError, AttributeError) as e:
+                    logger.warning(
+                        f"Verifier issue: {e}, "
+                        f"retrying... ({retries + 1}/{max_retries})"
+                    )
+                    retries += 1
+                    continue
 
-        if validate:
-            for i, item in enumerate(self._raw_data):
                 try:
-                    # Use DataPoint for validation only
-                    dp = DataPoint(**item)
-                    self.data.append(dp.to_dict())
+                    new_datapoint = DataPoint(
+                        question=agent_output["question"],
+                        rationale=rationale,
+                        final_answer=verifier_response.result,
+                        metadata={
+                            "synthetic": str(True),
+                            "created": datetime.now().isoformat(),
+                        },
+                    )
                 except ValidationError as e:
-                    logger.error(f"Sample {i} validation error: {e}")
-                    raise ValueError(f"Sample {i} validation error: {e}")
-        else:
-            # Skip validation and just store the data dictionaries
-            self.data = [dict(item) for item in self._raw_data]
-
-    def __getitem__(self, index: int) -> T_co:
-        r"""Get an item from the dataset.
+                    logger.warning(
+                        f"Datapoint validation failed: {e}, "
+                        f"retrying... ({retries + 1}/{max_retries})"
+                    )
+                    retries += 1
+                    continue
 
-        Args:
-            index (int): Index of the item to get.
+                valid_data_points.append(new_datapoint)
 
-        Returns:
-            T_co: Item from the dataset, possibly transformed.
-        """
-        sample = self.data[index]
+            except Exception as e:
+                logger.warning(
+                    f"Unexpected error: {e}, retrying..."
+                    f" ({retries + 1}/{max_retries})"
+                )
+                retries += 1
 
-        # Apply transformations if provided
-        if self.transform is not None:
-            sample = self.transform(sample)
+        if len(valid_data_points) < n:
+            raise RuntimeError(
+                f"Failed to generate {n} valid datapoints "
+                f"after {max_retries} retries."
+            )
 
-        return sample  # type: ignore[return-value]
+        self._data.extend(valid_data_points)
+        return valid_data_points
 
     def __len__(self) -> int:
-        r"""Return the size of the dataset.
-
-        Returns:
-            int: Number of samples in the dataset.
-        """
-        return len(self.data)
+        r"""Return the size of the dataset."""
+        return len(self._data)
 
-    @classmethod
-    def from_datapoints(
-        cls, datapoints: List[DataPoint], **kwargs
-    ) -> 'PyTorchDataset':
-        r"""Create a PyTorchDataset from a list of DataPoints.
+    def __getitem__(self, idx: int) -> DataPoint:
+        r"""Retrieve a datapoint by index.
 
         Args:
-            datapoints (List[DataPoint]): List of DataPoint objects.
-            **kwargs: Additional arguments to pass to the constructor.
+            idx (int): Index of the datapoint.
 
         Returns:
-            PyTorchDataset: A new PyTorchDataset instance.
-        """
-        data = [dp.to_dict() for dp in datapoints]
-        # We can skip validation since datapoints are already validated
-        return cls(data, validate=False, **kwargs)
-
-    def to_hf_dataset(self) -> HFDataset:
-        r"""Convert to a HuggingFace dataset.
-
-        Returns:
-            HFDataset: Dataset in HuggingFace format.
-        """
-        return HFDataset.from_list(self.data)
+            DataPoint: The datapoint corresponding to the given index.
 
-    def save_to_disk(self, path: str) -> None:
-        r"""Save the dataset to disk using PyTorch.
-
-        Args:
-            path (str): Path to save the dataset to.
+        Raises:
+            IndexError: If idx is out of bounds.
         """
-        torch.save(self.data, path)
-
-    @classmethod
-    def load_from_disk(
-        cls,
-        path: str,
-        transform: Optional[Callable] = None,
-        target_transform: Optional[Callable] = None,
-    ) -> 'PyTorchDataset':
-        r"""Load a dataset from disk.
-
-        Args:
-            path (str): Path to load the dataset from.
-            transform (Optional[Callable]): Transform to apply to samples.
-                (default: :obj:`None`)
-            target_transform (Optional[Callable]): Transform to apply to
-                targets. (default: :obj:`None`)
+        if idx < 0 or idx >= len(self._data):
+            raise IndexError(
+                f"Index {idx} out of bounds for dataset of "
+                f"size {len(self._data)}"
+            )
+        return self._data[idx]
 
-        Returns:
-            PyTorchDataset: Loaded dataset.
-        """
-        data = torch.load(path)
-        return cls(
-            data,
-            transform=transform,
-            target_transform=target_transform,
-            validate=False,
-        )
+    def save_to_jsonl(self, file_path: Union[str, Path]) -> None:
+        r"""Saves the dataset to a JSONL (JSON Lines) file.
 
-    @staticmethod
-    def collate_fn(
-        batch: List[Dict[str, Any]],
-    ) -> Dict[str, Union[List[Any], torch.Tensor]]:
-        r"""Collate function for PyTorch DataLoader.
+        Each datapoint is stored as a separate JSON object on a new line.
 
         Args:
-            batch (List[Dict[str, Any]]): Batch of samples from the dataset.
+            file_path (Union[str, Path]): Path to save the JSONL file.
 
-        Returns:
-            Dict[str, Union[List[Any], torch.Tensor]]: Collated batch with
-                tensors for numerical data.
-        """
-        if not batch:
-            return {}
-
-        # Initialize result dictionary with keys from first item - start with
-        # lists only
-        result: Dict[str, List[Any]] = {k: [] for k in batch[0].keys()}
-
-        # Collect values by key
-        for item in batch:
-            for k, v in item.items():
-                result[k].append(v)
-
-        # Convert numeric/boolean lists to tensors where possible
-        result_with_tensors: Dict[str, Union[List[Any], torch.Tensor]] = {}
-        for k, v in result.items():
-            if all(isinstance(x, (int, float, bool)) for x in v):
-                try:
-                    result_with_tensors[k] = torch.tensor(v)
-                except (ValueError, TypeError):
-                    # Keep as list if tensor conversion fails
-                    result_with_tensors[k] = v
-            else:
-                result_with_tensors[k] = v
-
-        return result_with_tensors
-
-    def get_dataloader(
-        self,
-        batch_size: int = 32,
-        shuffle: bool = True,
-        num_workers: int = 0,
-        pin_memory: bool = False,
-        **kwargs,
-    ) -> "DataLoader":
-        r"""Create a PyTorch DataLoader for this dataset.
-
-        Args:
-            batch_size (int): Batch size. (default: :obj:`32`)
-            shuffle (bool): Whether to shuffle the dataset. (default:
-                :obj:`True`)
-            num_workers (int): Number of workers for data loading. (default:
-            :obj:`0`)
-            pin_memory (bool): Whether to pin memory for faster GPU transfer.
-                (default: :obj:`False`)
-            **kwargs: Additional arguments to pass to DataLoader.
+        Raises:
+            ValueError: If the dataset is empty.
+            IOError: If there is an issue writing to the file.
 
-        Returns:
-            torch.utils.data.DataLoader: DataLoader for this dataset.
+        Notes:
+            - Uses `self._data`, which contains the generated datapoints.
+            - Overwrites the file if it already exists.
+            - Ensures compatibility with large datasets by using JSONL format.
         """
-        from torch.utils.data import DataLoader
-
-        return DataLoader(
-            self,
-            batch_size=batch_size,
-            shuffle=shuffle,
-            num_workers=num_workers,
-            collate_fn=self.collate_fn,
-            pin_memory=pin_memory,
-            **kwargs,
-        )
-
-
-def convert_hf_to_pytorch(
-    hf_dataset: HFDataset,
-    transform: Optional[Callable] = None,
-    target_transform: Optional[Callable] = None,
-    column_mapping: Optional[Dict[str, str]] = None,
-    validate: bool = True,
-    batch_size: Optional[int] = None,
-) -> Union["PyTorchDataset", "DataLoader"]:
-    r"""Convert a HuggingFace dataset to a PyTorchDataset or DataLoader.
-
-    This function maps HuggingFace dataset columns to the expected DataPoint
-    format, validates the data, and creates a PyTorchDataset or DataLoader.
-
-    Args:
-        hf_dataset (HFDataset): HuggingFace dataset to convert.
-        transform (Optional[Callable]): Transform to apply to samples.
-        target_transform (Optional[Callable]): Transform to apply to targets.
-        column_mapping (Optional[Dict[str, str]]): Mapping from HuggingFace
-            column names to DataPoint field names. If None, assumes columns
-            already match DataPoint fields.
-        validate (bool): Whether to validate data points using DataPoint
-            schema. (default: :obj:`True`)
-        batch_size (Optional[int]): If provided, returns a DataLoader with the
-            specified batch size instead of a PyTorchDataset. (default:
-            :obj:`None`)
-
-    Returns:
-        Union[PyTorchDataset, torch.utils.data.DataLoader]: Converted dataset
-            or DataLoader if batch_size is provided.
-    """
-    # Convert HuggingFace dataset to list of dicts more efficiently
-    mapped_dataset = []
-
-    for i in range(len(hf_dataset)):
-        item = hf_dataset[i]
-        if column_mapping is not None:
-            # Apply column mapping if provided
-            mapped_item = {}
-            for hf_col, dp_field in column_mapping.items():
-                if hf_col in item:
-                    mapped_item[dp_field] = item[hf_col]
-            mapped_dataset.append(mapped_item)
-        else:
-            # Otherwise use item directly
-            mapped_dataset.append(dict(item))
-
-    # Create PyTorchDataset
-    dataset: PyTorchDataset = PyTorchDataset(
-        mapped_dataset,
-        transform=transform,
-        target_transform=target_transform,
-        validate=validate,
-    )
-
-    # Return DataLoader if batch_size is provided
-    if batch_size is not None:
-        return dataset.get_dataloader(batch_size=batch_size)
+        if not self._data:
+            raise ValueError("Dataset is empty. No data to save.")
 
-    return dataset
+        file_path = Path(file_path)
 
-
-def convert_synthetic_to_pytorch(
-    synthetic_dataset: 'SyntheticDataset',
-    transform: Optional[Callable] = None,
-    target_transform: Optional[Callable] = None,
-    batch_size: Optional[int] = None,
-) -> Union["PyTorchDataset", "DataLoader"]:
-    r"""Convert a SyntheticDataset to a PyTorchDataset or DataLoader.
-
-    Args:
-        synthetic_dataset (SyntheticDataset): Synthetic dataset to convert.
-        transform (Optional[Callable]): Transform to apply to samples.
-        target_transform (Optional[Callable]): Transform to apply to targets.
-        batch_size (Optional[int]): If provided, returns a DataLoader with the
-            specified batch size instead of a PyTorchDataset. (default:
-            :obj:`None`)
-
-    Returns:
-        Union[PyTorchDataset, torch.utils.data.DataLoader]: Converted dataset
-            or DataLoader if batch_size is provided.
-    """
-    dataset = PyTorchDataset.from_datapoints(
-        synthetic_dataset.data,
-        transform=transform,
-        target_transform=target_transform,
-    )
-
-    # Return DataLoader if batch_size is provided
-    if batch_size is not None:
-        return dataset.get_dataloader(batch_size=batch_size)
-
-    return dataset
-
-
-def save_synthetic_dataset(
-    synthetic_dataset: 'SyntheticDataset',
-    path: str,
-    compression: bool = True,
-) -> None:
-    r"""Save a synthetic dataset to disk using PyTorch format.
-
-    Args:
-        synthetic_dataset (SyntheticDataset): Dataset to save.
-        path (str): Path to save the dataset to.
-        compression (bool): Whether to use compression to reduce file size.
-            (default: :obj:`True`)
-    """
-    pytorch_dataset = convert_synthetic_to_pytorch(synthetic_dataset)
-
-    # Save with compression if enabled (uses less disk space)
-    if compression:
-        torch.save(
-            pytorch_dataset.data,  # type: ignore[union-attr]
-            path,
-            _use_new_zipfile_serialization=True,
-        )
-    else:
-        pytorch_dataset.save_to_disk(path)  # type: ignore[union-attr]
-
-
-def load_pytorch_dataset(
-    path: str,
-    transform: Optional[Callable] = None,
-    target_transform: Optional[Callable] = None,
-    batch_size: Optional[int] = None,
-) -> Union["PyTorchDataset", "DataLoader"]:
-    r"""Load a PyTorchDataset from disk.
-
-    Args:
-        path (str): Path to load the dataset from.
-        transform (Optional[Callable]): Transform to apply to samples.
-        target_transform (Optional[Callable]): Transform to apply to targets.
-        batch_size (Optional[int]): If provided, returns a DataLoader with the
-            specified batch size instead of a PyTorchDataset. (default:
-            :obj:`None`)
-
-    Returns:
-        Union[PyTorchDataset, torch.utils.data.DataLoader]: Loaded dataset or
-            DataLoader if batch_size is provided.
-    """
-    dataset = PyTorchDataset.load_from_disk(
-        path, transform=transform, target_transform=target_transform
-    )
-
-    # Return DataLoader if batch_size is provided
-    if batch_size is not None:
-        return dataset.get_dataloader(batch_size=batch_size)
-
-    return dataset
+        try:
+            with file_path.open("w", encoding="utf-8") as f:
+                for datapoint in self._data:
+                    json.dump(datapoint.to_dict(), f)
+                    f.write("\n")  # Ensure each entry is on a new line
+            logger.info(f"Dataset saved successfully to {file_path}")
+        except IOError as e:
+            logger.error(f"Error writing to file {file_path}: {e}")
+            raise