#!/usr/bin/env bash
set -euo pipefail

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
workdir="$repo_root/compare_model/LaCo"

model_path="${MODEL_PATH:-meta-llama/Llama-2-7b-hf}"
target_layers="${TARGET_LAYERS:-16}"
merge_layers="${MERGE_LAYERS:-2}"
interval="${INTERVAL:-1}"
lowest_layer="${LOWEST_LAYER:-0}"
threshold="${THRESHOLD:-0.45}"
dtype="${DTYPE:-bfloat16}"
device="${DEVICE:-cuda}"
max_prompt_length="${MAX_PROMPT_LENGTH:-128}"
output_dir="${OUTDIR:-$repo_root/results/laco_llama_target_${target_layers}}"

python_args=(
  --model_path "$model_path"
  --output_dir "$output_dir"
  --target_layers "$target_layers"
  --merge_layers "$merge_layers"
  --interval "$interval"
  --lowest_layer "$lowest_layer"
  --threshold "$threshold"
  --dtype "$dtype"
  --device "$device"
  --max_prompt_length "$max_prompt_length"
)

if [[ "${TRUST_REMOTE_CODE:-0}" == "1" ]]; then
  python_args+=(--trust_remote_code)
fi

if [[ "${FORCE_TARGET:-1}" == "1" ]]; then
  python_args+=(--force_target)
else
  python_args+=(--no_force_target)
fi

if [[ -n "${PROMPT_FILE:-}" ]]; then
  python_args+=(--prompt_file "$PROMPT_FILE")
fi

if [[ -n "${SAVE_LAYERS:-}" ]]; then
  # shellcheck disable=SC2206
  save_layers=(${SAVE_LAYERS})
  python_args+=(--save_layers "${save_layers[@]}")
fi

python_args+=("$@")

mkdir -p "$output_dir"
git_commit="unknown"
if git -C "$repo_root" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
  git_commit=$(git -C "$repo_root" rev-parse HEAD)
fi
{
  echo "git_commit=$git_commit"
  echo "command:"
  printf '%q ' python "$repo_root/compare_model/LaCo/laco_llama.py" "${python_args[@]}"
  echo
} > "$output_dir/run_args.txt"

cd "$workdir"
PYTHONPATH="$repo_root${PYTHONPATH:+:$PYTHONPATH}" \
python laco_llama.py "${python_args[@]}"