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

export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}

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)

MODEL_PATH=${MODEL_PATH:-/data/local/hujk/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/llama31-8b_tri_native_tool_response_only_chunked}

python "$SCRIPT_DIR/_tuning.fix.chunked.py" \
  --model_path "$MODEL_PATH" \
  --data_path "$DATA_PATH" \
  --head_path "$HEAD_PATH" \
  --output_dir "$OUTPUT_DIR" \
  --topk "${TOPK:-100p}" \
  --epochs "${EPOCHS:-30}" \
  --batch_size "${BATCH_SIZE:-6}" \
  --batch-save-interval "${BATCH_SAVE_INTERVAL:-100}" \
  --lr "${LR:-5e-4}" \
  --lambda_preserve "${LAMBDA_PRESERVE:-1.0}" \
  --attn-chunk-size "${ATTN_CHUNK_SIZE:-4096}" \
  --max-train-steps "${MAX_TRAIN_STEPS:--1}" \
  --repeat-single-sample "${REPEAT_SINGLE_SAMPLE:-1}" \
  --max-len "${MAX_LEN:-50000}" \
  --logits-to-keep "${LOGITS_TO_KEEP:-1}" \
  ${SKIP_SAVE_EVAL:+--skip-save-eval} \
  ${NO_SAVE:+--no-save} \
  ${GRADIENT_CHECKPOINTING:+--gradient-checkpointing}