#!/usr/bin/env sh
set -eu

export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
export PYTORCH_ALLOC_CONF=${PYTORCH_ALLOC_CONF:-expandable_segments:True}

CONDA_BIN=${CONDA_BIN:-}
if [ -z "$CONDA_BIN" ]; then
  if command -v conda >/dev/null 2>&1; then
    CONDA_BIN=$(command -v conda)
  elif [ -x /opt/miniconda/bin/conda ]; then
    CONDA_BIN=/opt/miniconda/bin/conda
  fi
fi

if [ -n "$CONDA_BIN" ]; then
  eval "$("$CONDA_BIN" shell.bash hook)"
  conda activate focallora4
fi

SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
BASE=$(CDPATH= cd -- "$SCRIPT_DIR/.." && pwd)
ROOT=$(CDPATH= cd -- "$BASE/.." && pwd)
export EVAL_MMLU=1
export EVAL_TOPICATTACK=1
MODEL_PATH=${MODEL_PATH:-$ROOT/models/Llama-3.1-8B-Instruct}
DATA_PATH=${DATA_PATH:-$BASE/3-1_model_training_data_gen/single_turn/tri_native_tool_response_only.json}
HEAD_PATH=${HEAD_PATH:-$BASE/2-2_head_identification_scoring/model_score/sep_Llama-3.1-8B-Instruct/heads_sorted/all_roc_inst_0.1.json}
OUTPUT_DIR=${OUTPUT_DIR:-$SCRIPT_DIR/outputs_lora/attn_kl_clean}

python "$SCRIPT_DIR/train_attn_kl_clean.py" \
  --model-path "$MODEL_PATH" \
  --data-path "$DATA_PATH" \
  --head-path "$HEAD_PATH" \
  --output-dir "$OUTPUT_DIR" \
  --topk "${TOPK:-18.75p}" \
  --epochs "${EPOCHS:-3}" \
  --batch-size "${BATCH_SIZE:-4}" \
  --max-train-steps "${MAX_TRAIN_STEPS:--1}" \
  --repeat-single-sample "${REPEAT_SINGLE_SAMPLE:-1}" \
  --max-len "${MAX_LEN:-50000}" \
  --lr "${LR:-1e-4}" \
  --attn-chunk-size "${ATTN_CHUNK_SIZE:-4096}" \
  --save-steps "${SAVE_STEPS:-100}" \
  --log-interval "${LOG_INTERVAL:-100}" \
  --log-csv "${LOG_CSV:-}" \
  --eval-output "${EVAL_OUTPUT:-}" \
  --mmlu-split "${MMLU_SPLIT:-dev}" \
  --mmlu-size "${MMLU_SIZE:-256}" \
  --mmlu-max-new-tokens "${MMLU_MAX_NEW_TOKENS:-16}" \
  --eval-size "${EVAL_SIZE:-24}" \
  --dev-holdout "${DEV_HOLDOUT:-0}" \
  --eval-data-path "${EVAL_DATA_PATH:-$BASE/1_raw_dataset/topicattack/data/crafted_instruction_data_squad_injection_qa.json}" \
  --eval-topicattack-path "${EVAL_TOPIC_PATH:-$BASE/1_raw_dataset/topicattack/data/crafted_instruction_data_squad_conversation_attack_complete.json}" \
  --eval-config "${EVAL_CONFIG:-native_tool_response_only}" \
  --eval-attacks "${EVAL_ATTACKS:-none,naive,ignore,escape_separation,completion_realcmb,conv_attack}" \
  --eval-attack-side "${EVAL_ATTACK_SIDE:-end}" \
  --eval-batch-size "${EVAL_BATCH_SIZE:-4}" \
  --eval-max-new-tokens "${EVAL_MAX_NEW_TOKENS:-256}" \
  ${GRADIENT_CHECKPOINTING:+--gradient-checkpointing} \
  ${EVAL_AFTER_TRAIN:+--eval-after-train} \
  ${EVAL_STEP0:+--eval-step0} \
  ${EVAL_MMLU:+--eval-mmlu} \
  ${EVAL_TOPICATTACK:+--eval-topicattack}