#!/usr/bin/env bash
# Run AFTER training stops (GPU free). Authoritative base-vs-merged comparison on identical
# sets: vuln-detection, pentest MCQ, MMLU-security, MMLU-general (forgetting), HumanEval (coding).
# Usage: bash finalize.sh <checkpoint-dir>
set -uo pipefail
cd /workspace/infosec
export HF_HOME=/workspace/hf-cache HF_HUB_CACHE=/workspace/hf-cache/hub PYTHONUNBUFFERED=1
CKPT="${1:?usage: finalize.sh <checkpoint-dir>}"
MERGED=/workspace/checkpoints/qwen36_a100_stage1_merged
EVALS="--eval data/eval/vuln_detection_test.jsonl --eval data/eval/knowledge_mcq.jsonl --eval data/eval/mmlu_security.jsonl --eval data/eval/mmlu_general.jsonl --report-dir reports/eval --sample 50 --max-new-tokens 256"

echo "== [1/5] eval BASE (knowledge/vuln) =="
python training/scripts/eval_hf_model.py --model Qwen/Qwen3.6-27B --label base_a100 $EVALS
echo "== [2/5] eval BASE (HumanEval coding) =="
python training/scripts/eval_coding.py --model Qwen/Qwen3.6-27B --label base_a100 --report-dir reports/eval

echo "== [3/5] merge checkpoint $CKPT =="
python training/scripts/merge_lora.py --config training/configs/stage1_a100.yaml --adapter "$CKPT"

echo "== [4/5] eval MERGED (knowledge/vuln) =="
python training/scripts/eval_hf_model.py --model "$MERGED" --label merged $EVALS
echo "== [5/5] eval MERGED (HumanEval coding) =="
python training/scripts/eval_coding.py --model "$MERGED" --label merged --report-dir reports/eval

echo "== COMPARISON (base vs merged) =="
python3 - <<'PY'
import json, os
def collect(label):
    out={}
    ev=f"reports/eval/{label}_eval.json"
    if os.path.exists(ev):
        for r in json.load(open(ev)).get("results",[]):
            name=os.path.basename(r.get("file",r.get("kind","?"))).replace(".jsonl","")
            if "accuracy" in r:
                c=f"acc={round(r['accuracy'],3)}"
                if r.get("kind")=="vuln_detection": c+=f" f1={round(r.get('f1_vuln',0),3)}"
                out[name]=c
            else: out[name]=r.get("error","?")
    cd=f"reports/eval/{label}_coding.json"
    if os.path.exists(cd):
        for r in json.load(open(cd)).get("results",[]):
            out["humaneval"]=f"pass@1={r.get('pass@1')}"
    return out
b=collect("base_a100"); m=collect("merged")
keys=sorted(set(b)|set(m))
print(f"{'benchmark':28}{'BASE':24}{'MERGED':24}")
for k in keys:
    print(f"{k:28}{str(b.get(k,'-')):24}{str(m.get(k,'-')):24}")
json.dump({"base":b,"merged":m}, open("reports/eval/comparison.json","w"), indent=2)
PY
echo FINALIZE_DONE