Update Phase 3 to next-token train-from-scratch with checkpoint auto-resume

examples/4_pretrain_from_scratch.py

@@ -13,10 +13,12 @@ from torch.utils.data import Dataset
 from pathlib import Path
 from transformers import (
     AutoTokenizer,
+    AutoConfig,
     AutoModelForMaskedLM,
     Trainer,
     TrainingArguments,
 )
+from transformers.trainer_utils import get_last_checkpoint
 
 
 class PretrainingDataset(Dataset):
@@ -96,7 +98,8 @@ def main():
     print("=" * 80)
 
     model_id = "mineself2016/GeneMamba"
-    checkpoint_dir = Path("./from_scratch_pretrain/checkpoint-last")
+    output_dir = "./from_scratch_pretrain"
+    checkpoint_dir = Path(output_dir) / "checkpoint-last"
     
     # ============================================================
     # Step 1: Load tokenizer spec
@@ -117,33 +120,29 @@ def main():
     # ============================================================
     print("\n[Step 2] Building model (resume if checkpoint exists)...")
 
-    from configuration_genemamba import GeneMambaConfig
-    from modeling_genemamba import GeneMambaForMaskedLM
-    
-    model_config = GeneMambaConfig(
-        vocab_size=25426,
-        hidden_size=256,  # Smaller for faster demo
-        num_hidden_layers=12,  # Reduced for demo
-        intermediate_size=1024,
-        max_position_embeddings=2048,
-        mamba_mode="mean",
-        embedding_pooling="mean",
-        num_labels=2,
-        hidden_dropout_prob=0.1,
-        initializer_range=0.02,
-    )
-    
+    model_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    model_config.vocab_size = 25426
+    model_config.hidden_size = 256
+    model_config.num_hidden_layers = 12
+    model_config.intermediate_size = 1024
+    model_config.max_position_embeddings = 2048
+    model_config.mamba_mode = "mean"
+
+    resume_from_checkpoint = None
     if checkpoint_dir.exists():
+        resume_from_checkpoint = str(checkpoint_dir)
+    else:
+        resume_from_checkpoint = get_last_checkpoint(output_dir)
+
+    if resume_from_checkpoint is not None:
         model = AutoModelForMaskedLM.from_pretrained(
-            str(checkpoint_dir),
+            resume_from_checkpoint,
             trust_remote_code=True,
             local_files_only=True,
         )
-        resume_from_checkpoint = str(checkpoint_dir)
-        print(f"✓ Found checkpoint, resume from: {checkpoint_dir}")
+        print(f"✓ Found checkpoint, resume from: {resume_from_checkpoint}")
     else:
-        model = GeneMambaForMaskedLM(model_config)
-        resume_from_checkpoint = None
+        model = AutoModelForMaskedLM.from_config(model_config, trust_remote_code=True)
         print("✓ No checkpoint found, start from scratch")
     
     # Count parameters
@@ -186,8 +185,6 @@ def main():
     # ============================================================
     print("\n[Step 5] Setting up training...")
     
-    output_dir = "./from_scratch_pretrain"
-    
     training_args = TrainingArguments(
         output_dir=output_dir,
         num_train_epochs=5,
@@ -297,8 +294,8 @@ def main():
     print("Phase 3 Complete! Model trained from scratch and ready to use.")
     print("=" * 80)
     
-    return model, trainer, config
+    return model, trainer, model_config
 
 
 if __name__ == "__main__":
-    model, trainer, config = main()
+    model, trainer, model_config = main()