Fix: support preformatted text datasets in train_eagle3.py (avoid forcing conversations generator)

scripts/train_eagle3.py

@@ -17,7 +17,7 @@
 from tqdm import tqdm
 from transformers import AutoProcessor, AutoTokenizer
 
-from datasets import Dataset
+from datasets import Dataset, load_dataset
 from specforge import (
     AutoDraftModelConfig,
     AutoEagle3DraftModel,
@@ -231,6 +231,20 @@ def parse_args() -> Tuple[ArgumentParser, Namespace]:
     return parser, args
 
 
+def _load_dataset_from_path(data_path: str, is_preformatted: bool):
+    if is_preformatted:
+        return load_dataset(
+            "json",
+            data_files=data_path,
+            split="train",
+        )
+    else:
+        return Dataset.from_generator(
+            generator=safe_conversations_generator,
+            gen_kwargs={"file_path": data_path},
+        )
+
+
 def build_tracker(args: Namespace, parser: ArgumentParser) -> Tracker:
     """
     Build the experiment tracker according to the report_to argument.
@@ -452,12 +466,10 @@ def build_dataloaders(
         f"{args.max_length}-"
         f"{args.chat_template}-"
         f"{args.target_model_path}"  # Tokenizer may also different
+        f"is_preformatted={args.is_preformatted}"
     )
     cache_key = hashlib.md5(cache_params_string.encode()).hexdigest()
-    train_dataset = Dataset.from_generator(
-        generator=safe_conversations_generator,
-        gen_kwargs={"file_path": args.train_data_path},
-    )
+    train_dataset = _load_dataset_from_path(args.train_data_path, args.is_preformatted)
     is_online = (
         args.train_data_path is not None and args.train_hidden_states_path is None
     )
@@ -505,9 +517,8 @@ def build_dataloaders(
     )
     if args.eval_data_path is not None or args.eval_hidden_states_path is not None:
         if args.eval_data_path is not None:
-            eval_dataset = Dataset.from_generator(
-                generator=safe_conversations_generator,
-                gen_kwargs={"file_path": args.eval_data_path},
+            eval_dataset = _load_dataset_from_path(
+                args.eval_data_path, args.is_preformatted
             )
             eval_eagle3_dataset = build_eagle3_dataset(
                 eval_dataset,