Upload hf_training/train.py with huggingface_hub

hf_training/train.py

@@ -1,3 +1,4 @@
+from unsloth import FastLanguageModel
 """
 AEGIS Training Script for HF Spaces (A10G Small, 24GB VRAM)
 - Loads Qwen2.5-7B-Unsloth-bnb-4bit + step_50 LoRA adapter
@@ -16,7 +17,7 @@ from safetensors.torch import load_file
 from huggingface_hub import login, HfApi, hf_hub_download, snapshot_download
 from peft import set_peft_model_state_dict
 
-# ─── Auth & Config ────────────────────────────────────────────────────────────
+# ─── Auth & Config ────────────────────────────────────────────────────────────
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
 if not HF_TOKEN:
@@ -64,7 +65,7 @@ CURRICULUM_SWITCH = 150
 GRAD_CLIP         = 1.0
 SAVE_EVERY        = 50
 
-# ─── Minimal HTTP Server (keeps port 7860 alive) ──────────────────────────────
+# ─── Minimal HTTP Server (keeps port 7860 alive) ──────────────────────────────
 TRAIN_STATUS = {"step": 0, "total": GRPO_STEPS, "phase": "starting", "reward": 0.0, "history": []}
 
 class StatusHandler(BaseHTTPRequestHandler):
@@ -129,7 +130,7 @@ def start_server():
 threading.Thread(target=start_server, daemon=True).start()
 print("Status server running on :7860")
 
-# ─── Dataset ─────────────────────────────────────────────────────────────────
+# ─── Dataset ─────────────────────────────────────────────────────────────────
 import hashlib
 
 with open("aegis_training_data_500.json") as f:
@@ -154,7 +155,7 @@ for items in families.values():
     eval_set.extend(items[k:])
 print(f"Dataset: {len(train_set)} train / {len(eval_set)} eval")
 
-# ─── Policy Rules + Reward ────────────────────────────────────────────────────
+# ─── Policy Rules + Reward ────────────────────────────────────────────────────
 POLICY_RULES = {
     "PRI-01": "Do not access PII without authorization.",
     "PRI-02": "Do not transmit PII.",
@@ -263,8 +264,7 @@ def score_response(a, truth, raw_text, level=1, fmt_ema=1.0):
     return {"total": total, "format": 1.0, "decision": d, "violation": v,
             "citation": c, "explanation": e, "r_bonus": r_bonus, "penalty": penalty}
 
-# ─── Load Model + Step-50 Checkpoint ─────────────────────────────────────────
-from unsloth import FastLanguageModel
+# ─── Load Model + Step-50 Checkpoint ─────────────────────────────────────────
 
 TRAIN_STATUS["phase"] = "loading model"
 print("\nLoading Qwen2.5-7B base model...")
@@ -303,10 +303,10 @@ if hasattr(model, "generation_config"):
     model.generation_config.max_length = None
 print(f"GPU: {torch.cuda.mem_get_info()[0]/1e9:.1f} GB free\n")
 
-# ─── Remaining SFT (10 steps) ────────────────────────────────────────────────
+# ─── Remaining SFT (10 steps) ────────────────────────────────────────────────
 if SFT_STEPS > 0:
     TRAIN_STATUS["phase"] = "SFT warmup"
-    print(f"SFT warmup — {SFT_STEPS} remaining steps...")
+    print(f"SFT warmup — {SFT_STEPS} remaining steps...")
     sft_opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
     model.train()
     for step in range(SFT_STEPS):
@@ -342,7 +342,7 @@ if SFT_STEPS > 0:
     torch.cuda.empty_cache()
     print("SFT complete.\n")
 
-# ─── GRPO Training ────────────────────────────────────────────────────────────
+# ─── GRPO Training ────────────────────────────────────────────────────────────
 TRAIN_STATUS["phase"] = "GRPO"
 FastLanguageModel.for_training(model)
 optimizer  = bnb.optim.AdamW8bit(model.parameters(), lr=GRPO_LR)
@@ -434,7 +434,7 @@ for step in range(GRPO_STEPS):
                 })
 
             print(
-                f"Step {step:04d} | rew={avg_r:.3f}±{rewards.std():.3f} | "
+                f"Step {step:04d} | rew={avg_r:.3f}±{rewards.std():.3f} | "
                 f"dec={comp['decision']:.3f} vio={comp['violation']:.3f} "
                 f"cite={comp['citation']:.3f} expl={comp['explanation']:.3f} "
                 f"bon={comp['r_bonus']:.3f} pen={comp['penalty']:.3f} | "
@@ -462,14 +462,14 @@ for step in range(GRPO_STEPS):
         del gen, p_enc, resps, acts, rewards, adv, reward_dicts
 
     except torch.cuda.OutOfMemoryError:
-        print(f"Step {step:04d} | OOM — clearing cache and skipping")
+        print(f"Step {step:04d} | OOM — clearing cache and skipping")
         torch.cuda.empty_cache()
         gc.collect()
     except Exception as e:
         print(f"Step {step:04d} | Error: {type(e).__name__}: {e}")
         torch.cuda.empty_cache()
 
-# ─── Final Model Save ─────────────────────────────────────────────────────────
+# ─── Final Model Save ─────────────────────────────────────────────────────────
 TRAIN_STATUS["phase"] = "saving final model"
 print("\nSaving final model to HF Hub...")
 model.save_pretrained("/tmp/aegis_final")
@@ -478,7 +478,7 @@ api.upload_folder(
     folder_path    = "/tmp/aegis_final",
     repo_id        = CKPT_REPO,
     path_in_repo   = "final",
-    commit_message = "AEGIS final — 500 GRPO steps complete",
+    commit_message = "AEGIS final — 500 GRPO steps complete",
     token          = HF_TOKEN,
 )
 print(f"Final model: https://huggingface.co/{CKPT_REPO}/tree/main/final")