fix: auto alpaca

dataquality/__init__.py

@@ -31,7 +31,7 @@
 """
 
 
-__version__ = "1.1.2"
+__version__ = "1.1.3"
 
 import sys
 from typing import Any, List, Optional

dataquality/dq_auto/base_data_manager.py

@@ -5,11 +5,6 @@
 from datasets import Dataset, DatasetDict, load_dataset
 
 from dataquality.exceptions import GalileoException
-from dataquality.integrations.seq2seq.formatter import (
-    BaseFormatter,
-    DefaultFormatter,
-    get_formatter,
-)
 from dataquality.schemas.split import Split
 from dataquality.utils.auto import (
     _apply_column_mapping,
@@ -20,9 +15,6 @@
 class BaseDatasetManager:
     DEMO_DATASETS: List[str] = []
 
-    def __init__(self) -> None:
-        self.formatter: BaseFormatter = DefaultFormatter()
-
     def _validate_dataset_dict(
         self,
         dd: DatasetDict,
@@ -148,16 +140,13 @@ def try_load_dataset_dict(
         if hf_data:
             if isinstance(hf_data, str):
                 dd = load_dataset(hf_data)
-                self.formatter = get_formatter(hf_data)
             else:
                 dd = hf_data
             assert isinstance(dd, DatasetDict), (
                 "hf_data must be a path to a huggingface DatasetDict in the hf hub or "
                 "a DatasetDict object. "
                 "If this is just a Dataset, pass it to `train_data`"
             )
-            # Apply the datasets custom formatter on load dataset dict
-            dd = dd.map(self.formatter.format_sample)
             return dd
 
         return None

dataquality/dq_auto/schema.py

@@ -4,6 +4,8 @@
 import pandas as pd
 from datasets import Dataset, DatasetDict
 
+from dataquality.integrations.seq2seq.formatter import BaseFormatter, DefaultFormatter
+
 
 @dataclass
 class BaseAutoDatasetConfig:
@@ -46,6 +48,9 @@ class BaseAutoDatasetConfig:
     # Column names
     input_col: str = "text"
     target_col: str = "label"
+    # Dataset input / output formatter
+    max_train_size: Optional[int] = None
+    formatter: BaseFormatter = DefaultFormatter()
 
     def __post_init__(self) -> None:
         if not any([self.hf_data, self.train_path, self.train_data]):

dataquality/integrations/seq2seq/auto.py

@@ -19,6 +19,7 @@
     add_val_data_if_missing,
     get_meta_cols,
     run_name_from_hf_dataset,
+    sample_dataset_dict,
 )
 from dataquality.utils.torch import cleanup_cuda
 
@@ -54,7 +55,7 @@ def try_load_dataset_dict_from_config(
             hf_data = dataset_config.hf_data
             if isinstance(hf_data, str):
                 dd = load_dataset(hf_data)
-                self.formatter = get_formatter(hf_data)
+                dataset_config.formatter = get_formatter(hf_data)
             elif isinstance(hf_data, DatasetDict):
                 dd = hf_data
             else:
@@ -64,8 +65,8 @@ def try_load_dataset_dict_from_config(
                     "If this is just a Dataset, pass it to `train_data`"
                 )
 
-            # Apply the datasets custom formatter on load dataset dict
-            dd = dd.map(self.formatter.format_sample)
+            dataset_config.input_col = dataset_config.formatter.input_col
+            dataset_config.target_col = dataset_config.formatter.target_col
             return dd, dataset_config
 
         return None, dataset_config
@@ -109,6 +110,9 @@ def get_dataset_dict_from_config(
             if test_data is not None:
                 dd[Split.test] = self._convert_to_hf_dataset(test_data)
 
+        # Apply the datasets custom formatter on load dataset dict
+        dd = dd.map(dataset_config.formatter.format_sample)
+        dd = sample_dataset_dict(dd, dataset_config)
         return self._validate_dataset_dict(dd, []), dataset_config
 
     def _validate_dataset_dict(
@@ -139,11 +143,15 @@ def _log_dataset_dict(dd: DatasetDict, input_col: str, target_col: str) -> None:
     for key in dd.keys():
         ds: Dataset = dd[key]
         if key in Split.get_valid_keys():
-            meta = get_meta_cols(ds.features, {input_col, target_col})
             if input_col != "text" and "text" in ds.column_names:
                 ds = ds.rename_columns({"text": "_metadata_text"})
             if target_col != "label" and "label" in ds.column_names:
                 ds = ds.rename_columns({"label": "_metadata_label"})
+            if input_col != "input" and "input" in ds.column_names:
+                ds = ds.rename_columns({"input": "_metadata_input"})
+            if target_col != "target" and "target" in ds.column_names:
+                ds = ds.rename_columns({"target": "_metadata_target"})
+            meta = get_meta_cols(ds.features, {input_col, target_col})
             dq.log_dataset(ds, text=input_col, label=target_col, split=key, meta=meta)
 
 
@@ -231,7 +239,6 @@ def auto(
     dq.init(TaskType.seq2seq, project_name=project_name, run_name=run_name)
     input_col = dataset_config.input_col
     target_col = dataset_config.target_col
-
     # We 'watch' in get_trainer, which must happen before logging datasets
     model, dataloaders = get_trainer(
         dd,

dataquality/integrations/seq2seq/formatter.py

@@ -1,13 +1,14 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Dict, Type
+from typing import Dict, Optional, Type
 
 
 @dataclass
 class BaseFormatter(ABC):
     name: str
     input_col: str
     target_col: str
+    max_train_size: Optional[int] = None
 
     @abstractmethod
     def format_sample(self, sample: Dict[str, str]) -> Dict[str, str]:
@@ -31,6 +32,7 @@ class AlpacaFormatter(BaseFormatter):
     name: str = "tatsu-lab/alpaca"
     input_col: str = "formatted_input"
     target_col: str = "output"
+    max_train_size: int = 1000
 
     def format_sample(self, sample: Dict[str, str]) -> Dict[str, str]:
         """Formats the alpaca dataset for seq2seq

dataquality/utils/auto.py

@@ -8,12 +8,46 @@
 import pandas as pd
 from datasets import ClassLabel, Dataset, DatasetDict, load_dataset
 
+from dataquality.dq_auto.schema import BaseAutoDatasetConfig
 from dataquality.exceptions import GalileoException, GalileoWarning
 from dataquality.schemas.split import Split
 from dataquality.schemas.task_type import TaskType
 from dataquality.utils.name import BAD_CHARS_REGEX
 
 
+def sample_dataset_dict(
+    dd: DatasetDict, dataset_config: BaseAutoDatasetConfig
+) -> DatasetDict:
+    """Samples the dataset dict to the max train size
+
+    A few important notes:
+    - If max train size is greater than the dataset size, we don't sample
+    - If max train size is None we also don't sample
+    - We set max eval size to be 25% of max train size
+    - Test and inference data are not sampled
+    """
+    max_train_sz = (
+        dataset_config.max_train_size or dataset_config.formatter.max_train_size
+    )
+    if not max_train_sz:
+        return dd
+
+    max_eval_sz = int(max_train_sz * 0.25)
+    for split, dataset in dd.items():
+        sampled_size = len(dataset)
+        if split == Split.train:
+            sampled_size = min(sampled_size, max_train_sz)
+        elif split == Split.validation:
+            sampled_size = min(sampled_size, max_eval_sz)
+
+        if len(dataset) > sampled_size:
+            # Slice the dataset to the max size
+            dataset = dataset.select(range(sampled_size))
+            dd[split] = dataset
+
+    return dd
+
+
 def get_meta_cols(
     cols: Iterable, reserved_cols: Optional[Set[str]] = None
 ) -> List[str]: