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

README.md

@@ -211,7 +211,8 @@ If a checkpoint exists, resume automatically; otherwise start from scratch.
 ```python
 import torch
 from pathlib import Path
-from transformers import AutoTokenizer, AutoModelForMaskedLM, Trainer, TrainingArguments
+from transformers import AutoTokenizer, AutoConfig, AutoModelForMaskedLM, Trainer, TrainingArguments
+from transformers.trainer_utils import get_last_checkpoint
 
 tokenizer = AutoTokenizer.from_pretrained(
     "mineself2016/GeneMamba",
@@ -222,30 +223,31 @@ print("vocab_size:", tokenizer.vocab_size)  # 25426
 print("unk/pad:", tokenizer.unk_token_id, tokenizer.pad_token_id)  # 0, 1
 print("cls/mask:", tokenizer.cls_token_id, tokenizer.mask_token_id)  # None, None
 
-# Build model config
-from configuration_genemamba import GeneMambaConfig
-from modeling_genemamba import GeneMambaForMaskedLM
-
-config = GeneMambaConfig(
-    vocab_size=25426,
-    hidden_size=512,
-    num_hidden_layers=24,
-    max_position_embeddings=2048,
-    mamba_mode="mean",
-)
+# Build model config (no local modeling file import required)
+config = AutoConfig.from_pretrained("mineself2016/GeneMamba", trust_remote_code=True)
+config.vocab_size = 25426
+config.hidden_size = 512
+config.num_hidden_layers = 24
+config.max_position_embeddings = 2048
+config.mamba_mode = "mean"
 
 # Resume if checkpoint exists
-checkpoint_dir = Path("./from_scratch_pretrain/checkpoint-last")
+output_dir = "./from_scratch_pretrain"
+checkpoint_dir = Path(output_dir) / "checkpoint-last"
+
 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)
 else:
-    model = GeneMambaForMaskedLM(config)
-    resume_from_checkpoint = None
+    model = AutoModelForMaskedLM.from_config(config, trust_remote_code=True)
 
 class NextTokenTrainer(Trainer):
     def compute_loss(self, model, inputs, return_outputs=False):
@@ -262,7 +264,7 @@ class NextTokenTrainer(Trainer):
 trainer = NextTokenTrainer(
     model=model,
     args=TrainingArguments(
-        output_dir="./from_scratch_pretrain",
+        output_dir=output_dir,
         num_train_epochs=3,
         per_device_train_batch_size=32,
         learning_rate=2e-5,