🆔 Add datast_config to ScriptArguments (#2440)

* datast_config_name * Update trl/utils.py [ci skip] * sort import * typo [ci skip] * Trigger CI * Rename `dataset_config_name` to `dataset_config`
huggingface · Dec 10, 2024 · 6a05fef · 6a05fef
1 parent 2f72f47
commit 6a05fef
Show file tree

Hide file tree

Showing 20 changed files with 27 additions and 19 deletions.
diff --git a/examples/scripts/bco.py b/examples/scripts/bco.py
@@ -126,7 +126,7 @@ def mean_pooling(model_output, attention_mask):
     if tokenizer.chat_template is None:
         model, tokenizer = setup_chat_format(model, tokenizer)
 
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     accelerator = Accelerator()
     embedding_model = AutoModel.from_pretrained(

diff --git a/examples/scripts/cpo.py b/examples/scripts/cpo.py
@@ -81,7 +81,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
 

diff --git a/examples/scripts/dpo.py b/examples/scripts/dpo.py
@@ -111,7 +111,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     ##########
     # Training

diff --git a/examples/scripts/dpo_online.py b/examples/scripts/dpo_online.py
@@ -121,7 +121,7 @@
     if tokenizer.pad_token_id is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     trainer = OnlineDPOTrainer(
         model=model,

diff --git a/examples/scripts/dpo_vlm.py b/examples/scripts/dpo_vlm.py
@@ -104,7 +104,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     ################
     # Training

diff --git a/examples/scripts/gkd.py b/examples/scripts/gkd.py
@@ -104,7 +104,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     with PartialState().local_main_process_first():
         dataset = dataset.map(

diff --git a/examples/scripts/kto.py b/examples/scripts/kto.py
@@ -91,7 +91,7 @@
         model, tokenizer = setup_chat_format(model, tokenizer)
 
     # Load the dataset
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     # Initialize the KTO trainer
     trainer = KTOTrainer(

diff --git a/examples/scripts/nash_md.py b/examples/scripts/nash_md.py
@@ -121,7 +121,7 @@
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
 
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     trainer = NashMDTrainer(
         model=model,

diff --git a/examples/scripts/orpo.py b/examples/scripts/orpo.py
@@ -81,7 +81,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
 

diff --git a/examples/scripts/ppo/ppo.py b/examples/scripts/ppo/ppo.py
@@ -119,7 +119,9 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name, split=script_args.dataset_train_split)
+    dataset = load_dataset(
+        script_args.dataset_name, name=script_args.dataset_config, split=script_args.dataset_train_split
+    )
     eval_samples = 100
     train_dataset = dataset.select(range(len(dataset) - eval_samples))
     eval_dataset = dataset.select(range(len(dataset) - eval_samples, len(dataset)))

diff --git a/examples/scripts/ppo/ppo_tldr.py b/examples/scripts/ppo/ppo_tldr.py
@@ -126,7 +126,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
     train_dataset = dataset[script_args.dataset_train_split]
     eval_dataset = dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None
 

diff --git a/examples/scripts/reward_modeling.py b/examples/scripts/reward_modeling.py
@@ -107,7 +107,7 @@
     ##############
     # Load dataset
     ##############
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     ##########
     # Training

diff --git a/examples/scripts/rloo/rloo.py b/examples/scripts/rloo/rloo.py
@@ -90,7 +90,9 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name, split=script_args.dataset_train_split)
+    dataset = load_dataset(
+        script_args.dataset_name, name=script_args.dataset_config, split=script_args.dataset_train_split
+    )
     eval_samples = 100
     train_dataset = dataset.select(range(len(dataset) - eval_samples))
     eval_dataset = dataset.select(range(len(dataset) - eval_samples, len(dataset)))

diff --git a/examples/scripts/rloo/rloo_tldr.py b/examples/scripts/rloo/rloo_tldr.py
@@ -92,7 +92,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
     train_dataset = dataset[script_args.dataset_train_split]
     eval_dataset = dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None
 

diff --git a/examples/scripts/sft.py b/examples/scripts/sft.py
@@ -90,7 +90,7 @@
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     ################
     # Training

diff --git a/examples/scripts/sft_video_llm.py b/examples/scripts/sft_video_llm.py
@@ -166,7 +166,7 @@ class CustomScriptArguments(ScriptArguments):
     training_args.dataset_kwargs = {"skip_prepare_dataset": True}
 
     # Load dataset
-    dataset = load_dataset(script_args.dataset_name, split="train")
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config, split="train")
 
     # Setup model
     torch_dtype = (

diff --git a/examples/scripts/sft_vlm.py b/examples/scripts/sft_vlm.py
@@ -109,7 +109,7 @@ def collate_fn(examples):
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     ################
     # Training

diff --git a/examples/scripts/sft_vlm_smol_vlm.py b/examples/scripts/sft_vlm_smol_vlm.py
@@ -121,7 +121,7 @@ def collate_fn(examples):
     ################
     # Dataset
     ################
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     ################
     # Training

diff --git a/examples/scripts/xpo.py b/examples/scripts/xpo.py
@@ -106,7 +106,7 @@
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
 
-    dataset = load_dataset(script_args.dataset_name)
+    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
 
     trainer = XPOTrainer(
         model=model,

diff --git a/trl/utils.py b/trl/utils.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
+from typing import Optional
 
 
 @dataclass
@@ -23,6 +24,8 @@ class ScriptArguments:
     Args:
         dataset_name (`str`):
             Dataset name.
+        dataset_config (`str` or `None`, *optional*, defaults to `None`):
+            Dataset configuration name. Corresponds to the `name` argument of the [`~datasets.load_dataset`] function.
         dataset_train_split (`str`, *optional*, defaults to `"train"`):
             Dataset split to use for training.
         dataset_test_split (`str`, *optional*, defaults to `"test"`):
@@ -35,6 +38,7 @@ class ScriptArguments:
     """
 
     dataset_name: str
+    dataset_config: Optional[str] = None
     dataset_train_split: str = "train"
     dataset_test_split: str = "test"
     gradient_checkpointing_use_reentrant: bool = False