Update train.py

train.py

@@ -20,7 +20,7 @@ from utils import ensure_dirs, setup_logger
 
 
 # ==============================
-# 🔥 FIXED BACKUP CALLBACK
+# 🔥 BACKUP CALLBACK
 # ==============================
 class BackupCallback(TrainerCallback):
     def on_save(self, args, state, control, **kwargs):
@@ -46,42 +46,15 @@ class BackupCallback(TrainerCallback):
 
             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
-    return (path / "config.json").exists()
+    return path.exists() and (path / "config.json").exists()
 
 
 def _resolve_model_path(logger) -> Path:
@@ -93,8 +66,7 @@ def _resolve_model_path(logger) -> Path:
 
     if _is_valid_hf_model_dir(fallback):
         logger.warning(
-            "Primary model path %s is missing HF files. Falling back to %s",
-            primary.resolve(),
+            "Primary model missing → using fallback %s",
             fallback.resolve(),
         )
         return fallback
@@ -103,7 +75,7 @@ def _resolve_model_path(logger) -> Path:
 
 
 # ==============================
-# BUILD MODEL
+# BUILD MODEL (FIXED)
 # ==============================
 def _build_model_and_tokenizer(model_path: Path):
     tokenizer = AutoTokenizer.from_pretrained(
@@ -115,13 +87,14 @@ def _build_model_and_tokenizer(model_path: Path):
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
+    # 🔥 FIXED MODEL LOADING
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
         trust_remote_code=True,
-        local_files_only=True,
-        dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        use_safetensors=True,   # IMPORTANT
     )
 
+    # LoRA
     lora_cfg = LoraConfig(
         r=16,
         lora_alpha=32,
@@ -136,7 +109,7 @@ def _build_model_and_tokenizer(model_path: Path):
 
 
 # ==============================
-# SMART RESUME
+# CHECKPOINT RESUME (SAFE)
 # ==============================
 def get_latest_checkpoint(checkpoint_dir):
     if not os.path.exists(checkpoint_dir):
@@ -162,19 +135,18 @@ 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}")
+        logger.info(f"Trying resume from: {latest_checkpoint}")
         try:
             trainer.train(resume_from_checkpoint=latest_checkpoint)
             return
         except Exception as e:
-            logger.warning(f"Resume failed: {e}")
+            logger.warning(f"Resume failed → starting fresh: {e}")
 
-    logger.warning("No valid checkpoint → starting fresh training")
     trainer.train()
 
 
 # ==============================
-# MAIN TRAIN FUNCTION
+# MAIN TRAIN
 # ==============================
 def train(resume: bool):
     ensure_dirs([
@@ -210,7 +182,6 @@ def train(resume: bool):
         logging_steps=50,
         save_steps=250,
         save_total_limit=3,
-        gradient_checkpointing=False,
         report_to="none",
         remove_unused_columns=False,
     )
@@ -229,11 +200,11 @@ def train(resume: bool):
     trainer.model.save_pretrained(str(PATHS.lora_output_dir))
     tokenizer.save_pretrained(str(PATHS.tokenizer_output_dir))
 
-    print("\n✅ Training complete. Model saved.")
+    print("\n✅ Training complete.")
 
 
 # ==============================
-# ENTRY POINT
+# ENTRY
 # ==============================
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()