prepare dataset

scripts/prepare_contrain_dataset.py

@@ -67,6 +67,21 @@ def batch_iterator(name=None):
         del dataset
         gc.collect()
 
+    if name in (None, 'datatab/ultrachat_200k_serbian'):
+        dataset = load_dataset('datatab/ultrachat_200k_serbian', split='train')
+
+        for row in dataset:
+            text = [
+                f"<|im_start|>{n['role']}\n{n['content']}<|im_end|>"
+                for n in row['messages_srb']
+            ]
+
+            text = '\n'.join(text) + '\n'
+            yield text
+
+        del dataset
+        gc.collect()
+
     if name in (None, 'datatab/ultrafeedback_binarized_serbian'):
         dataset = load_dataset('datatab/ultrafeedback_binarized_serbian', split='train_sft')
 
@@ -112,6 +127,44 @@ def batch_iterator(name=None):
         del dataset
         gc.collect()
 
+    if name in (None, 'datatab/orca_math_world_problem_200k_serbian'):
+        dataset = load_dataset('datatab/orca_math_world_problem_200k_serbian', split='train')
+
+        for row in dataset:
+            text = []
+
+            text.append(
+                '<|im_start|>user\n'
+                f"{row['question_translated_srb']}<|im_end|>"
+            )
+
+            text.append(
+                '<|im_start|>assistant\n'
+                f"{row['answer_translated_srb']}<|im_end|>"
+            )
+
+            text = '\n'.join(text) + '\n'
+            yield text
+
+        del dataset
+        gc.collect()
+
+    if name in (None, 'datatab/open-orca-slim-serbian'):
+        dataset = load_dataset('datatab/open-orca-slim-serbian', split='train')
+        role_map = {'system': 'system', 'human': 'user', 'gpt': 'assistant'}
+
+        for row in dataset['conversations']:
+            text = [
+                f"<|im_start|>{role_map[n['role']]}\n{n['value']}<|im_end|>"
+                for n in row['chosen']
+            ]
+
+            text = '\n'.join(text) + '\n'
+            yield text
+
+        del dataset
+        gc.collect()
+
 
 def tokenize_fn(dataset_name, tokenizer=None):
     for text in batch_iterator(dataset_name):
@@ -123,8 +176,11 @@ datasets_names = [
     'Replete-AI/Everything_Instruct_Multilingual',
     'HuggingFaceH4/ultrachat_200k',
     'HuggingFaceH4/no_robots',
+    'datatab/ultrachat_200k_serbian',
     'datatab/ultrafeedback_binarized_serbian',
     'datatab/alpaca-cleaned-serbian-full',
+    'datatab/orca_math_world_problem_200k_serbian',
+    'datatab/open-orca-slim-serbian',
 ]
 
 outputs = optimize(