sft trainer auto concats multiple datasets

lucidrains · Feb 1, 2024 · bded2cc · bded2cc
1 parent e9a582c
commit bded2cc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py b/self_rewarding_lm_pytorch/self_rewarding_lm_pytorch.py
@@ -623,7 +623,7 @@ class FinetuneConfig:
 
 @dataclass
 class SFTConfig(FinetuneConfig):
-    train_dataset: Dataset
+    train_dataset: Union[Dataset, List[Dataset]]
     valid_dataset: Optional[Dataset] = None
     dropout: float = 0.1
     trainer_kwargs: dict = field(default_factory = dict)
@@ -677,7 +677,7 @@ class SelfPlayConfig(FinetuneConfig):
 @beartype
 def create_default_paper_config(
     *,
-    train_sft_dataset: Dataset,
+    train_sft_dataset: Union[Dataset, List[Dataset]],
     self_reward_prompt_dataset: Union[Dataset, Tuple[Dataset, Dataset]],
     valid_sft_dataset: Optional[Dataset] = None,
     num_generated_preference_pairs = (3964, 6942),

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'self-rewarding-lm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.1',
+  version = '0.2.2',
   license='MIT',
   description = 'Self Rewarding LM - Pytorch',
   author = 'Phil Wang',