Adding continued_pretraining task (#131)

* add continued pretraining script

* simplify config; add dataset_config option

* add ds configs in data mixer creator

* use extended sftconfig

* add option to avoid setting chat template

* fix data_configs bug

* add continued pretraining info

* add gpt2-nl recipe for continued pretraining example

* add final newline

* make style

* Update README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* Update README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* Update recipes/gpt2-nl/README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* rename continued pretraining to cpt

* improve README

---------

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

src/alignment/configs.py

@@ -57,6 +57,7 @@ class H4ArgumentParser(HfArgumentParser):
             inputs = {k: v for k, v in vars(data_yaml).items() if k in keys}
             for arg, val in other_args.items():
                 # add only if in keys
+
                 if arg in keys:
                     base_type = data_yaml.__dataclass_fields__[arg].type
                     inputs[arg] = val
@@ -201,10 +202,18 @@ class DataArguments:
         default=None,
         metadata={"help": ("Datasets and their proportions to be used for training ift/rl.")},
     )
+    text_column: Optional[str] = field(
+        default="text",
+        metadata={"help": "The column name to use for the text in the dataset (only used for continued pretraining)."},
+    )
     dataset_splits: Optional[List[str]] = field(
         default_factory=lambda: ["train", "test"],
         metadata={"help": ("List of train test splits to use in the dataset")},
     )
+    dataset_configs: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": "List of dataset config names. If given must be the same length as 'dataset_mixer' keys."},
+    )
     preprocessing_num_workers: Optional[int] = field(
         default=None,
         metadata={"help": "The number of processes to use for the preprocessing."},
@@ -226,6 +235,7 @@ class DataArguments:
 class SFTConfig(transformers.TrainingArguments):
     """
     Arguments related to the training process itself. For all parameters, see: https://huggingface.co/docs/transformers/v4.26.1/en/main_classes/trainer#transformers.TrainingArguments
+    Also used for the continued pretraining task.
     """
 
     dataset_kwargs: Optional[Dict[str, Any]] = field(

src/alignment/data.py

@@ -98,6 +98,7 @@ def apply_chat_template(
 def get_datasets(
     data_config: DataArguments | dict,
     splits: List[str] = ["train", "test"],
+    configs: Optional[List[str]] = None,
     shuffle: bool = True,
 ) -> DatasetDict:
     """
@@ -133,32 +134,40 @@ def get_datasets(
     else:
         raise ValueError(f"Data config {data_config} not recognized.")
 
-    raw_datasets = mix_datasets(dataset_mixer, splits=splits, shuffle=shuffle)
+    raw_datasets = mix_datasets(dataset_mixer, splits=splits, configs=configs, shuffle=shuffle)
     return raw_datasets
 
 
-def mix_datasets(dataset_mixer: dict, splits: Optional[List[str]] = None, shuffle=True) -> DatasetDict:
+def mix_datasets(
+    dataset_mixer: dict, configs: Optional[List[str]] = None, splits: Optional[List[str]] = None, shuffle=True
+) -> DatasetDict:
     """
     Loads and mixes datasets according to proportions specified in `dataset_mixer`.
 
     Args:
         dataset_mixer (`dict`):
             Dictionary containing the dataset names and their training proportions. By default, all test proportions are 1.
+        configs (Optional[List[str]], *optional*, defaults to `None`):
+            List of dataset config names. If given must be the same length as 'dataset_mixer' keys.
         splits (Optional[List[str]], *optional*, defaults to `None`):
             Dataset splits to load and mix. Assumes the splits exist in all datasets and have a `train_` or `test_` prefix.
         shuffle (`bool`, *optional*, defaults to `True`):
             Whether to shuffle the training and testing/validation data.
     """
+    configs = [None] * len(dataset_mixer) if not configs else configs
+    if configs is not None and len(configs) != len(dataset_mixer):
+        raise ValueError("The number of given dataset config names must be the same as the given number of datasets.")
+
     raw_datasets = DatasetDict()
     raw_train_datasets = []
     raw_val_datasets = []
     fracs = []
-    for ds, frac in dataset_mixer.items():
+    for (ds, frac), ds_config in zip(dataset_mixer.items(), configs):
         fracs.append(frac)
         for split in splits:
             try:
                 # Try first if dataset on a Hub repo
-                dataset = load_dataset(ds, split=split)
+                dataset = load_dataset(ds, ds_config, split=split)
             except DatasetGenerationError:
                 # If not, check local dataset
                 dataset = load_from_disk(os.path.join(ds, split))

src/alignment/model_utils.py

@@ -62,7 +62,9 @@ def get_quantization_config(model_args: ModelArguments) -> BitsAndBytesConfig |
     return quantization_config
 
 
-def get_tokenizer(model_args: ModelArguments, data_args: DataArguments) -> PreTrainedTokenizer:
+def get_tokenizer(
+    model_args: ModelArguments, data_args: DataArguments, auto_set_chat_template: bool = True
+) -> PreTrainedTokenizer:
     """Get the tokenizer for the model."""
     tokenizer = AutoTokenizer.from_pretrained(
         model_args.model_name_or_path
@@ -82,7 +84,7 @@ def get_tokenizer(model_args: ModelArguments, data_args: DataArguments) -> PreTr
 
     if data_args.chat_template is not None:
         tokenizer.chat_template = data_args.chat_template
-    elif tokenizer.chat_template is None and tokenizer.default_chat_template is None:
+    elif auto_set_chat_template and tokenizer.chat_template is None and tokenizer.default_chat_template is None:
         tokenizer.chat_template = DEFAULT_CHAT_TEMPLATE
 
     return tokenizer