Update train.py

train.py

@@ -83,7 +83,24 @@ def format_example(ex):
 
 
 def prepare_dataset(tokenizer, dataset_name):
-    dataset = load_dataset(dataset_name, split="train")
+    """
+    Supports:
+    - gsm8k
+    - gsm8k:main
+    - any_dataset
+    """
+
+    # Auto-fix gsm8k without config
+    if dataset_name == "gsm8k":
+        dataset_name = "gsm8k:main"
+
+    # Handle dataset:config format
+    if ":" in dataset_name:
+        name, config = dataset_name.split(":", 1)
+        dataset = load_dataset(name, config, split="train")
+    else:
+        dataset = load_dataset(dataset_name, split="train")
+
     dataset = dataset.map(format_example, remove_columns=dataset.column_names)
 
     def tokenize(ex):