#!/bin/bash
#PBS -N p7_ar_finish
#PBS -q A100
#PBS -l select=1:ncpus=32:mem=120gb
#PBS -l walltime=02:00:00
#PBS -j oe
#PBS -o /home/bbabajan/nanochat-ar/project7_finalize_sft_eval.log

set -euo pipefail

ROOT=/home/bbabajan/nanochat-ar
NANO=$ROOT/nanochat
export NANOCHAT_BASE_DIR=$ROOT/nanochat_cache
export NANOCHAT_DTYPE=bfloat16
export OMP_NUM_THREADS=1
export WANDB_MODE=offline
export TORCHDYNAMO_DISABLE=1
export TORCH_COMPILE_DISABLE=1
export CUDA_MODULE_LOADING=LAZY
export HF_DATASETS_OFFLINE=1
export TRANSFORMERS_OFFLINE=1

cd "$ROOT"
source "$NANO/.venv_project7/bin/activate"

echo "PROJECT7_FINALIZE_START $(date -u +%Y-%m-%dT%H:%M:%SZ)"
echo "HOST=$(hostname)"
echo "NANOCHAT_BASE_DIR=$NANOCHAT_BASE_DIR"
python - <<'PY'
import os, torch
print("torch", torch.__version__)
print("cuda_available", torch.cuda.is_available())
print("cuda_device", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none")
print("base_dir", os.environ.get("NANOCHAT_BASE_DIR"))
PY

python scripts/07_make_arabic_sft.py \
  --input "$ROOT/data_clean/train.txt" \
  --output "$ROOT/data/arabic_sft_generated.jsonl" \
  --max-examples 2000

python scripts/07_patch_chat_sft_custom.py

cd "$NANO"
export PROJECT7_CUSTOM_SFT_JSONL="$ROOT/data/arabic_sft_generated.jsonl"

echo "PROJECT7_BASE_EVAL_START $(date -u +%Y-%m-%dT%H:%M:%SZ)"
python -m scripts.base_eval \
  --eval bpb,sample \
  --model-tag ar-d12-eager \
  --step 50 \
  --device-batch-size 4 \
  --split-tokens 131072
echo "PROJECT7_BASE_EVAL_DONE $(date -u +%Y-%m-%dT%H:%M:%SZ)"

echo "PROJECT7_SFT_START $(date -u +%Y-%m-%dT%H:%M:%SZ)"
python -m scripts.chat_sft \
  --run dummy \
  --model-tag ar-d12-eager \
  --model-step 50 \
  --load-optimizer 0 \
  --num-iterations 30 \
  --max-seq-len 1024 \
  --device-batch-size 2 \
  --total-batch-size 32768 \
  --eval-every 15 \
  --eval-tokens 32768 \
  --chatcore-every -1 \
  --mmlu-epochs 0 \
  --gsm8k-epochs 0
echo "PROJECT7_SFT_DONE $(date -u +%Y-%m-%dT%H:%M:%SZ)"

echo "PROJECT7_FINAL_OUTPUTS"
find "$NANOCHAT_BASE_DIR" -maxdepth 3 -type f \( -name 'model_*.pt' -o -name 'meta_*.json' -o -name '*.md' \) -printf "%p|%s|%TY-%Tm-%Td %TH:%TM\n" | sort
echo "PROJECT7_FINALIZE_DONE $(date -u +%Y-%m-%dT%H:%M:%SZ)"