[feature] added translation, rallio instruct tuning dataset, prosocial for safety, new summary dataset

model/supervised_finetuning/tests/test_datasets.py

@@ -7,10 +7,11 @@ from custom_datasets.dialogue_collator import DialogueDataCollator
 def test_all_datasets():
     qa_base = QA_DATASETS
     summarize_base = SUMMARIZATION_DATASETS
-    others = ["prompt_dialogue", "webgpt", "soda", "joke"]
+    others = ["prompt_dialogue", "webgpt", "soda", "joke", "instruct_tuning"]
+    translation = ["dive_mt", "wmt2019_zh-en", "wmt2019_ru-en", "wmt2019_de-en", "ted_trans_de-ja", "ted_trans_nl-en"]
 
     config = Namespace(cache_dir=".cache")
-    for dataset_name in others + qa_base + summarize_base:
+    for dataset_name in translation:
         print(dataset_name)
         train, eval = get_one_dataset(config, dataset_name)
         # sanity check
@@ -51,4 +52,4 @@ def test_collate_fn():
 
 
 if __name__ == "__main__":
-    test_collate_fn()
+    test_all_datasets()