#!/bin/bash
#SBATCH -p ds4ai
#SBATCH -c 16
#SBATCH -t 24:00:00
#SBATCH --mem=32gb
#SBATCH --gres=gpu:1
#SBATCH --output=logs/test_%j.out
#SBATCH --error=logs/test_%j.err

# Set HF_HOME to a scratch directory
export HF_HOME=/scratch/tgeng_lab/sun/init_weights

# ---- user params ----
PATH_TO_DATA="openai/gsm8k"
PATH_TO_MODEL="Qwen/Qwen2-1.5B-Instruct"
rho=0.2
lr=0.03
solver_prompt_idx=1

CUDA_VISIBLE_DEVICES=0 \
WANDB_RUN_NAME="LSeek-${PATH_TO_DATA//\//-}-${PATH_TO_MODEL//\//-}-rho${rho}-lr${lr}" \
python -u main.py \
  --dataset "$PATH_TO_DATA" \
  --model_name_or_path "$PATH_TO_MODEL" \
  --output_dir ./output \
  --start_data_idx 0 \
  --end_data_idx 1319 \
  --k $rho \
  --lr $lr \
  --solver_prompt_idx $solver_prompt_idx \
  --device cuda
