#!/bin/bash
set -Eeuo pipefail

# ITFormer 4GPU SFT + eval with the same training configuration as
# run_full_train_4gpu.sh, except the Qwen LLM is loaded with FlashAttention-2.
# Stage A is disabled so the existing TimeSeriesEncoder checkpoint is reused.

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
cd "$PROJECT_DIR"

RUN_ID="${RUN_ID:-$(date +%Y%m%d_%H%M%S)}"

export RUN_ID
export ADAPTER_TYPE="${ADAPTER_TYPE:-itformer}"
export RUN_STAGE_A="${RUN_STAGE_A:-false}"
export TS_ENCODER_CHECKPOINT="${TS_ENCODER_CHECKPOINT:-save/pretrain/model.safetensors}"
export LLM_ATTN_IMPLEMENTATION="${LLM_ATTN_IMPLEMENTATION:-flash_attention_2}"
export SFT_OUTPUT_DIR="${SFT_OUTPUT_DIR:-save/sft_itformer_flash_attn_qwen2.5_3B_${RUN_ID}}"
export LOG_FILE="${LOG_FILE:-logs/full_train_itformer_flash_attn_4gpu_${RUN_ID}.log}"

echo "Run id: $RUN_ID"
echo "Adapter type: $ADAPTER_TYPE"
echo "RUN_STAGE_A: $RUN_STAGE_A"
echo "TS encoder checkpoint: $TS_ENCODER_CHECKPOINT"
echo "LLM attention implementation: $LLM_ATTN_IMPLEMENTATION"
echo "SFT output dir: $SFT_OUTPUT_DIR"
echo "Log file: $LOG_FILE"

bash "${PROJECT_DIR}/scripts/run_full_train_4gpu.sh"