File size: 4,364 Bytes
76ba348 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | #!/usr/bin/env bash
# Resume the GRPO curriculum from stage 3.
# Pulls SFT and stage2_multi adapters from the Hub, then runs the
# remaining stages (3 and 4), eval, bake-demo, and final push.
#
# steps_per_stage is intentionally lowered to 150 (vs 200 elsewhere)
# because we want to finish before the next /data wipe.
set -euo pipefail
export HF_HUB_ENABLE_HF_TRANSFER=0
echo "[1/6] Installing GPU stack ..."
pip install -q --upgrade pip setuptools wheel
pip install -q "future>=1.0.0"
pip install -q "torch==2.10.0" "torchvision==0.25.0" \
--index-url https://download.pytorch.org/whl/cu128
pip install -q "transformers>=4.46,<5"
pip install -q --no-deps "unsloth @ git+https://github.com/unslothai/unsloth.git"
pip install -q --no-deps "unsloth_zoo"
pip install -q --no-deps "trl>=0.18.2,<=0.24.0" peft accelerate bitsandbytes
pip install -q "datasets>=3.4.1,<4.4.0" tyro tensorboard matplotlib sentencepiece protobuf huggingface_hub
pip install -q hf_transfer msgspec "torchao>=0.13.0" cut_cross_entropy || true
pip install -q -r requirements.txt
echo "[2/6] Building / verifying datasets ..."
python -m train.make_sft_dataset --n 600 --out data/sft_train.jsonl
python -m eval.make_holdout --out data/holdout.jsonl
echo "[3/6] Skipped (SFT + Stages 1,2 already on Hub)."
echo "[3.5/6] Pulling SFT and Stage 2 adapters from the Hub ..."
python - <<'PY'
import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
from huggingface_hub import snapshot_download
sft = snapshot_download(
"shivam2k3/opensoc-defender-grpo-sft",
local_dir="checkpoints/defender_sft_adapter",
local_dir_use_symlinks=False,
)
print("SFT adapter ->", sft)
stage2 = snapshot_download(
"shivam2k3/opensoc-defender-grpo-stage2_multi",
local_dir="checkpoints/defender_grpo/stage2_multi/adapter",
local_dir_use_symlinks=False,
)
print("Stage 2 adapter ->", stage2)
PY
echo "[4/6] Resuming GRPO from stage 3 (2 remaining stages, 150 steps each) ..."
python -m train.train_grpo \
--model unsloth/Qwen2.5-3B-Instruct \
--sft-adapter checkpoints/defender_grpo/stage2_multi/adapter \
--stages stage3_mixed,stage4_adversarial \
--steps-per-stage 150 --num-generations 8 \
--batch-size 2 --grad-accum 4 --lr 5e-6 \
--report-to tensorboard \
--out checkpoints/defender_grpo
echo "[5/6] Eval + plots ..."
python -m eval.eval \
--baseline unsloth/Qwen2.5-3B-Instruct \
--trained-adapter checkpoints/defender_grpo/stage4_adversarial/adapter \
--holdout data/holdout.jsonl --out-dir eval/results
python -m eval.plot_results --in eval/results/summary.json --out-dir eval/results
python -m eval.plot_training --grpo-root checkpoints/defender_grpo --out-dir eval/results
echo "[6/6] Baking demo data for the Gradio /demo Space ..."
python -m eval.bake_demo \
--baseline unsloth/Qwen2.5-3B-Instruct \
--trained-adapter checkpoints/defender_grpo/stage4_adversarial/adapter \
--n 50 --out data/demo_examples.json
if [ -n "${HF_TOKEN:-}" ] && [ -n "${HF_PUSH_TARGET:-}" ]; then
echo "[7/7] Uploading artifacts back to ${HF_PUSH_TARGET} ..."
python - <<'PY'
import os
from huggingface_hub import HfApi, upload_folder
token = os.environ["HF_TOKEN"]
target = os.environ["HF_PUSH_TARGET"]
adapter_repo = os.environ.get(
"HF_ADAPTER_REPO", target.split("/")[0] + "/opensoc-defender-grpo"
)
api = HfApi(token=token)
adapter_dir = "checkpoints/defender_grpo/stage4_adversarial/adapter"
if os.path.isdir(adapter_dir):
api.create_repo(adapter_repo, exist_ok=True, private=False)
upload_folder(repo_id=adapter_repo, folder_path=adapter_dir,
commit_message="GRPO-trained Qwen2.5-3B-Instruct LoRA defender adapter",
token=token)
print(" adapter ->", adapter_repo)
for p in [
"data/demo_examples.json",
"eval/results/summary.json",
"eval/results/bar_macro_f1.png",
"eval/results/bar_dismiss_on_malicious.png",
"eval/results/confusion_baseline_zero_shot.png",
"eval/results/confusion_opensoc_grpo.png",
"eval/results/training_curves.png",
"eval/results/training_kl_loss.png",
]:
if os.path.exists(p):
api.upload_file(
path_or_fileobj=p, path_in_repo=p,
repo_id=target,
commit_message="trained: refresh " + os.path.basename(p),
token=token,
)
print(" ", p, "->", target)
print("PIPELINE_DONE")
PY
fi
|