#!/usr/bin/env bash
# Run AFTER training stops (GPU free). Evaluates BASE on the A100 eval sets, merges the
# chosen LoRA checkpoint, and evaluates the MERGED model on the same sets -> authoritative
# base-vs-merged comparison (identical sets, same sample).
# 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>}"
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/3] eval BASE on A100 eval sets =="
python training/scripts/eval_hf_model.py --model Qwen/Qwen3.6-27B --label base_a100 $EVALS

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

echo "== [3/3] eval MERGED =="
python training/scripts/eval_hf_model.py --model /workspace/checkpoints/qwen36_a100_stage1_merged --label merged $EVALS

echo "== comparison (base vs merged, per eval set) =="
python3 - <<'PY'
import json, os
def by_file(p):
    try:
        d = json.load(open(p))
    except Exception as e:
        return {"_err": str(e)}
    out = {}
    for r in d.get("results", []):
        name = os.path.basename(r.get("file", r.get("kind", "?"))).replace(".jsonl", "")
        if "accuracy" in r:
            cell = f"acc={round(r['accuracy'],3)}"
            if r.get("kind") == "vuln_detection":
                cell += f" f1={round(r.get('f1_vuln',0),3)}"
            out[name] = cell
        else:
            out[name] = r.get("error", "?")
    return out
b = by_file("reports/eval/base_a100_eval.json")
m = by_file("reports/eval/merged_eval.json")
keys = sorted(set(b) | set(m))
print(f"{'eval set':28} {'BASE':22} {'MERGED':22}")
for k in keys:
    print(f"{k:28} {str(b.get(k,'-')):22} {str(m.get(k,'-')):22}")
PY
echo FINALIZE_DONE