#!/bin/bash
set -uo pipefail
export HF_HOME=/scratch/hf PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
while ! grep -q "TRAIN_DONE" /root/lug_s1.log 2>/dev/null; do
  if grep -qE "Traceback|OutOfMemoryError" /root/lug_s1.log; then echo "ERREUR: stage 1 en echec"; exit 1; fi
  sleep 60
done
echo "=== stage 1 fini ($(date +%H:%M:%S)) ==="
grep -A9 "EVAL FINALE" /root/lug_s1.log | head -10
/root/venv/bin/python /root/train_wbert.py --lang lug --out /scratch/runs/lug_s2 \
  --init /scratch/runs/lug_s1/best \
  --train /scratch/prep/manifests/waxal_lug_train.jsonl \
  --eval /scratch/prep/manifests/waxal_lug_validation.jsonl \
  --lr 1e-5 --epochs 10 --bs 8 --grad_accum 4 --num_workers 18 --eval_steps 100 > /root/lug_s2.log 2>&1
if ! grep -q "TRAIN_DONE" /root/lug_s2.log; then echo "ERREUR: stage 2 en echec"; exit 1; fi
echo "=== stage 2 fini ($(date +%H:%M:%S)) ==="
grep -A9 "EVAL FINALE" /root/lug_s2.log | head -10
mkdir -p /root/models
cp -r /scratch/runs/lug_s1/best /root/models/lug_s1_best
cp -r /scratch/runs/lug_s2/best /root/models/lug_s2_best
cp /scratch/runs/lug_s1/final_metrics.json /root/models/lug_s1_metrics.json 2>/dev/null
cp /scratch/runs/lug_s2/final_metrics.json /root/models/lug_s2_metrics.json 2>/dev/null
echo "=== checkpoints archives sur disque persistant ==="
/root/venv/bin/python /root/infer_submit.py --out /scratch/submissions/sub_probe_lug.csv \
  --model lug=/scratch/runs/lug_s2/best --score_val > /root/submit_probe.log 2>&1
if ! grep -q "SUBMIT_DONE" /root/submit_probe.log; then echo "ERREUR: soumission en echec"; exit 1; fi
grep -E "VAL lug|SOUMISSION" /root/submit_probe.log
cp /scratch/submissions/sub_probe_lug.csv /root/models/
echo "CHAIN_DONE"