Backup update: add latest checkpoints and train.py changes

backup_step8000.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13c3e5c401a567493b92bf02f6d4040f5b6f578c4c413b33362a0009d7405237
+size 84689731

backup_step8250.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dff8ef0900eeed1141b8aac59e1c45697ff3c804e4e2792568f4fdf5754e021
+size 84688227

backup_step8500.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d35f24763676911a2f605ff63e56a62f521bde805757d51b2e356a004d479e2e
+size 84695943

backup_step8750.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:007068138f8a165ff5a3fea9ed096a94bdf620d0007b013d8834d69bfc650628
+size 84696682

backup_step9000.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a69b305a69b77ea66f9feeaaaa3bbd7c4a08f7111bbd6cdd3b90e2e59a5b2e7b
+size 84704097

backup_step9250.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8724eceedfd4f8c4f87a14f1fa8c2019bcbfe9af6165e57aac020bb04c65fd5
+size 84699876

backup_step9500.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a728fcf9931e37ae37a3db4044170a254473aa08f9a10e958ce88987f2575d8c
+size 84705286

backup_step9750.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe8bbef08bb3ee21de186753bce613d4b050b4011d85378737d464e190db65a7
+size 84703357

train.py

@@ -1,6 +1,7 @@
 import argparse
 from pathlib import Path
-from typing import List
+import os
+import subprocess
 
 import torch
 from peft import LoraConfig, TaskType, get_peft_model
@@ -9,20 +10,78 @@ from transformers import (
     AutoTokenizer,
     Trainer,
     TrainingArguments,
+    TrainerCallback,
     set_seed,
 )
 
 from config import PATHS, TRAINING_CONFIG
-from dataset import LocalJsonlInstructionDataset, format_prompt
+from dataset import LocalJsonlInstructionDataset
 from utils import ensure_dirs, setup_logger
 
 
+# ==============================
+# 🔥 FIXED BACKUP CALLBACK
+# ==============================
+class BackupCallback(TrainerCallback):
+    def on_save(self, args, state, control, **kwargs):
+        try:
+            checkpoint_dir = os.path.join(
+                args.output_dir,
+                f"checkpoint-{state.global_step}"
+            )
+
+            if not os.path.exists(checkpoint_dir):
+                return
+
+            os.makedirs("backups", exist_ok=True)
+
+            backup_name = f"backup_step{state.global_step}.tar.gz"
+            backup_path = os.path.join("backups", backup_name)
+
+            print(f"\n[BACKUP] Creating backup for step {state.global_step}...")
+
+            subprocess.run([
+                "tar", "-czf", backup_path, checkpoint_dir
+            ], check=True)
+
+            print(f"[BACKUP] Saved: {backup_path}")
+
+            # =========================
+            # 🔥 FIXED NUMERIC SORT
+            # =========================
+            backups = [
+                f for f in os.listdir("backups")
+                if f.endswith(".tar.gz")
+            ]
+
+            backups = sorted(
+                backups,
+                key=lambda x: int(x.split("step")[1].split(".")[0])
+            )
+
+            # =========================
+            # KEEP LAST 5 BACKUPS
+            # =========================
+            if len(backups) > 5:
+                old_backup = backups[0]
+                old_path = os.path.join("backups", old_backup)
+
+                if os.path.isfile(old_path):
+                    os.remove(old_path)
+                    print(f"[BACKUP] Removed old backup: {old_backup}")
+
+        except Exception as e:
+            print(f"[BACKUP ERROR] {e}")
+            # Never crash training
+
+
+# ==============================
+# MODEL PATH RESOLUTION
+# ==============================
 def _is_valid_hf_model_dir(path: Path) -> bool:
     if not path.exists():
         return False
-    has_config = (path / "config.json").exists()
-    has_weights = (path / "model.safetensors").exists() or (path / "pytorch_model.bin").exists()
-    return has_config and has_weights
+    return (path / "config.json").exists()
 
 
 def _resolve_model_path(logger) -> Path:
@@ -31,6 +90,7 @@ def _resolve_model_path(logger) -> Path:
 
     if _is_valid_hf_model_dir(primary):
         return primary
+
     if _is_valid_hf_model_dir(fallback):
         logger.warning(
             "Primary model path %s is missing HF files. Falling back to %s",
@@ -38,36 +98,19 @@ def _resolve_model_path(logger) -> Path:
             fallback.resolve(),
         )
         return fallback
-    raise FileNotFoundError(
-        "No valid HuggingFace model directory found.\n"
-        f"Checked: {primary.resolve()} and {fallback.resolve()}.\n"
-        "Expected files: config.json + model.safetensors (or pytorch_model.bin)."
-    )
 
+    raise FileNotFoundError("No valid model directory found.")
 
-def _build_model_and_tokenizer(model_path: Path, logger):
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_path,
-            trust_remote_code=True,
-            local_files_only=True,
-            use_fast=True,
-        )
-    except Exception as fast_exc:
-        logger.warning("Fast tokenizer load failed: %s. Retrying with slow tokenizer.", fast_exc)
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_path,
-                trust_remote_code=True,
-                local_files_only=True,
-                use_fast=False,
-            )
-        except Exception as slow_exc:
-            raise RuntimeError(
-                "Tokenizer loading failed for both fast and slow modes. "
-                "Ensure tokenizer files exist in the model folder and install "
-                "`sentencepiece` (and optionally `tiktoken`) if required."
-            ) from slow_exc
+
+# ==============================
+# BUILD MODEL
+# ==============================
+def _build_model_and_tokenizer(model_path: Path):
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+        local_files_only=True,
+    )
 
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -87,84 +130,75 @@ def _build_model_and_tokenizer(model_path: Path, logger):
         task_type=TaskType.CAUSAL_LM,
         target_modules="all-linear",
     )
+
     model = get_peft_model(model, lora_cfg)
     return model, tokenizer
 
 
-def _maybe_resume_train(trainer: Trainer, logger, resume_requested: bool) -> None:
-    if not resume_requested:
-        trainer.train()
-        return
-
-    try:
-        trainer.train(resume_from_checkpoint=True)
-    except (ValueError, OSError) as exc:
-        logger.warning(
-            "Resume requested but no valid checkpoint found (%s). Starting fresh training.",
-            exc,
-        )
-        trainer.train()
+# ==============================
+# SMART RESUME
+# ==============================
+def get_latest_checkpoint(checkpoint_dir):
+    if not os.path.exists(checkpoint_dir):
+        return None
 
+    checkpoints = [
+        d for d in os.listdir(checkpoint_dir)
+        if d.startswith("checkpoint-")
+    ]
 
-def _generate_predictions(model, tokenizer, prompts: List[str], logger) -> None:
-    model.eval()
-    device = model.device
-    logger.info("Running post-training evaluation prompts.")
+    if not checkpoints:
+        return None
 
-    for prompt in prompts:
-        full_prompt = format_prompt(
-            instruction=prompt,
-            input_text="",
-            output_text="",
-        )
-        inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=TRAINING_CONFIG.eval_max_new_tokens,
-                do_sample=True,
-                temperature=0.2,
-                top_p=0.95,
-                pad_token_id=tokenizer.pad_token_id,
-            )
-        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        print("\n" + "=" * 80)
-        print(f"PROMPT: {prompt}")
-        print("-" * 80)
-        print(decoded)
-
-
-def train(resume: bool) -> Path:
-    ensure_dirs(
-        [
-            PATHS.data_dir,
-            PATHS.output_dir,
-            PATHS.logs_dir,
-            PATHS.checkpoint_dir,
-            PATHS.lora_output_dir,
-            PATHS.tokenizer_output_dir,
-        ]
+    checkpoints = sorted(
+        checkpoints,
+        key=lambda x: int(x.split("-")[-1])
     )
+
+    return os.path.join(checkpoint_dir, checkpoints[-1])
+
+
+def safe_train(trainer, checkpoint_dir, logger):
+    latest_checkpoint = get_latest_checkpoint(checkpoint_dir)
+
+    if latest_checkpoint:
+        logger.info(f"Resuming from checkpoint: {latest_checkpoint}")
+        try:
+            trainer.train(resume_from_checkpoint=latest_checkpoint)
+            return
+        except Exception as e:
+            logger.warning(f"Resume failed: {e}")
+
+    logger.warning("No valid checkpoint → starting fresh training")
+    trainer.train()
+
+
+# ==============================
+# MAIN TRAIN FUNCTION
+# ==============================
+def train(resume: bool):
+    ensure_dirs([
+        PATHS.data_dir,
+        PATHS.output_dir,
+        PATHS.logs_dir,
+        PATHS.checkpoint_dir,
+        PATHS.lora_output_dir,
+        PATHS.tokenizer_output_dir,
+    ])
+
     logger = setup_logger("train", PATHS.logs_dir / "train.log")
     set_seed(42)
-    if not torch.cuda.is_available():
-        logger.warning(
-            "CUDA is not available. Training will run on CPU, which is very slow and can limit practical model quality."
-        )
-
-    if not PATHS.train_jsonl.exists():
-        raise FileNotFoundError(
-            f"Training dataset not found: {PATHS.train_jsonl.resolve()}. "
-            "Run data_fetch.py first."
-        )
 
     model_path = _resolve_model_path(logger)
-    logger.info("Loading model and tokenizer from %s", model_path.resolve())
-    model, tokenizer = _build_model_and_tokenizer(model_path, logger)
+    logger.info("Loading model from %s", model_path)
+
+    model, tokenizer = _build_model_and_tokenizer(model_path)
     model.print_trainable_parameters()
 
-    train_dataset = LocalJsonlInstructionDataset(tokenizer, max_length=TRAINING_CONFIG.max_length)
-    logger.info("Loaded %d samples from %s", len(train_dataset), PATHS.train_jsonl.resolve())
+    train_dataset = LocalJsonlInstructionDataset(
+        tokenizer,
+        max_length=TRAINING_CONFIG.max_length
+    )
 
     training_args = TrainingArguments(
         output_dir=str(PATHS.checkpoint_dir),
@@ -173,56 +207,37 @@ def train(resume: bool) -> Path:
         gradient_accumulation_steps=TRAINING_CONFIG.gradient_accumulation_steps,
         learning_rate=TRAINING_CONFIG.learning_rate,
         fp16=torch.cuda.is_available(),
-        lr_scheduler_type="cosine",
-        warmup_ratio=0.03,
-        weight_decay=0.01,
-        max_grad_norm=1.0,
-        gradient_checkpointing=True,
-        group_by_length=True,
-        logging_steps=TRAINING_CONFIG.logging_steps,
-        save_steps=TRAINING_CONFIG.save_steps,
-        save_total_limit=4,
+        logging_steps=50,
+        save_steps=250,
+        save_total_limit=3,
+        gradient_checkpointing=False,
         report_to="none",
         remove_unused_columns=False,
-        dataloader_num_workers=2,
-        dataloader_pin_memory=torch.cuda.is_available(),
     )
 
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset,
+        callbacks=[BackupCallback()],
     )
 
-    logger.info("Starting training. Resume mode: %s", resume)
-    _maybe_resume_train(trainer, logger, resume_requested=resume)
+    logger.info("Starting training...")
+
+    safe_train(trainer, str(PATHS.checkpoint_dir), logger)
 
-    logger.info("Saving LoRA adapters to %s", PATHS.lora_output_dir.resolve())
     trainer.model.save_pretrained(str(PATHS.lora_output_dir))
     tokenizer.save_pretrained(str(PATHS.tokenizer_output_dir))
 
-    prompts = [
-        "Write a Python binary search function",
-        "Fix this Python bug: list index out of range",
-        "Create a FastAPI endpoint",
-    ]
-    _generate_predictions(model, tokenizer, prompts, logger)
-
-    print(f"\nLoRA adapters saved to: {PATHS.lora_output_dir.resolve()}")
-    print(f"Tokenizer saved to: {PATHS.tokenizer_output_dir.resolve()}")
-    return PATHS.lora_output_dir
-
-
-def _build_arg_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="LoRA fine-tuning for MINDI Python coding tasks.")
-    parser.add_argument(
-        "--no-resume",
-        action="store_true",
-        help="Disable automatic resume_from_checkpoint=True behavior.",
-    )
-    return parser
+    print("\n✅ Training complete. Model saved.")
 
 
+# ==============================
+# ENTRY POINT
+# ==============================
 if __name__ == "__main__":
-    args = _build_arg_parser().parse_args()
-    train(resume=not args.no_resume and TRAINING_CONFIG.resume_training)
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--no-resume", action="store_true")
+    args = parser.parse_args()
+
+    train(resume=not args.no_resume)