#!/usr/bin/env bash
# Per-CoT-step multi-objective experiment. Resumes from GRPO ckpt, replaces
# the global full-y InfoNCE with per-slot InfoNCE (each slot k → chunk k of
# the gold y). Capacity-diagnostic + TF + AR ablations on the final ckpt,
# then push to HF under per_slot_exp/.
set -uo pipefail

REPO="LauraGG/blt-reasoner-pilot1"
OUT="/home/ubuntu/work/blt_per_slot"
CFG="/home/ubuntu/experiments/blt_reasoner/configs/exp7b_per_slot.json"
RESUME_FROM="/home/ubuntu/work/blt_grpo_opt13/final"
LOG="/home/ubuntu/work/queue_per_slot.log"

log() { echo "[$(date +%T)] $*" | tee -a "$LOG"; }

mkdir -p "$OUT"
cd /home/ubuntu
export TOKENIZERS_PARALLELISM=false TRANSFORMERS_NO_ADVISORY_WARNINGS=1 HF_HUB_DISABLE_PROGRESS_BARS=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

log "==========================================="
log "PER-SLOT multi-objective experiment from $RESUME_FROM"
log "==========================================="
python3 -u -m experiments.blt_reasoner.train --config "$CFG" \
    --resume_from "$RESUME_FROM" \
    > "$OUT/train.log" 2>&1
log "train exit=$?"

log "Capacity diagnostic on per_slot final"
python3 -u -m experiments.blt_reasoner.scripts.capacity_diagnostic \
    --ckpt "$OUT/final" --config "$CFG" --n 100 --K 16 --max_new_tokens 128 \
    --out "$OUT/final/capacity_diagnostic.json" \
    > "$OUT/capacity.log" 2>&1
log "capacity diag exit=$?"

log "TF ablation"
python3 -u -m experiments.blt_reasoner.scripts.ablate_teacher_forced \
    --ckpt "$OUT/final" --config "$CFG" --n 200 --K 16 \
    --out "$OUT/final/ablation_teacher_forced.json" \
    > "$OUT/tf_eval.log" 2>&1
log "TF ablate exit=$?"

log "AR ablation"
python3 -u -m experiments.blt_reasoner.eval \
    --ckpt "$OUT/final" --config "$CFG" --n 200 --K 16 \
    --max_new_tokens 192 --temperature 0.0 \
    --out "$OUT/final/ablation_n200_K16.json" \
    > "$OUT/ar_eval.log" 2>&1
log "AR ablate exit=$?"

log "pushing per_slot_exp/ to HF"
python3 - <<PYEND
import os
from huggingface_hub import HfApi
token = os.environ.get("BLT_HF_TOKEN", "").strip()
assert token.startswith("hf_"), "BLT_HF_TOKEN missing"
api = HfApi(token=token)
api.upload_folder(folder_path="$OUT", path_in_repo="per_slot_exp",
                   repo_id="$REPO", repo_type="model",
                   commit_message="Per-CoT-step multi-objective experiment")
print("[push] done")
PYEND
log "queue_per_slot.sh DONE"