use concat instead of interleave

google · Jun 30, 2024 · 884ac1e · 884ac1e
1 parent f975891
commit 884ac1e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 12 deletions.
diff --git a/DiarizationLM/unsloth/config.py b/DiarizationLM/unsloth/config.py
@@ -1,6 +1,6 @@
 # DataPrep
 TRAINING_INPUT = {
-    "FISHER": ("/YOUT_DATA_PATH/FISHER_ENGLISH_TRAIN_FULL.json", 1.0),
+    "FISHER": ("/YOUT_DATA_PATH/FISHER_ENGLISH_TRAIN_FULL.json", 1),
 }
 EMIT_INPUT_LENGTH = 6000
 EMIT_TARGET_LENGTH = 6000

diff --git a/DiarizationLM/unsloth/dataprep.py b/DiarizationLM/unsloth/dataprep.py
@@ -1,5 +1,5 @@
 import config
-from datasets import Dataset, disable_caching, interleave_datasets
+from datasets import Dataset, disable_caching, concatenate_datasets
 from diarizationlm import utils
 
 
@@ -33,22 +33,16 @@ def build_dataset_single_source(input_file: str):
   )
   dataset1 = Dataset.from_generator(reader_hyp2ora.generate_data_dict)
   dataset2 = Dataset.from_generator(reader_deg2ref.generate_data_dict)
-  return [dataset1, dataset2]
+  return concatenate_datasets([dataset1, dataset2])
 
 
 def build_dataset():
   disable_caching()
   all_datasets = []
-  all_probs = []
   for data_name in config.TRAINING_INPUT:
-    data_path, data_prob = config.TRAINING_INPUT[data_name]
-    all_datasets.extend(build_dataset_single_source(data_path))
-    all_probs.extend([data_prob] * 2)
-  prob_sum = sum(all_probs)
-  all_probs = [prob / prob_sum for prob in all_probs]
-  dataset = interleave_datasets(all_datasets,
-                                probabilities=all_probs,
-                                stopping_strategy="all_exhausted")
+    data_path, data_repeat = config.TRAINING_INPUT[data_name]
+    all_datasets.extend([build_dataset_single_source(data_path)] * data_repeat)
+  dataset = concatenate_datasets(all_datasets)
   dataset = dataset.shuffle(seed=42)
   dataset = dataset.map(formatting_prompts_func)
   return dataset