#!/bin/bash
# This script automates the setup and training process for the HRM fine-tuning.
# The 'set -e' command ensures that the script will exit immediately if any command fails.
set -e

echo "--- Starting Automated Finetuning Workflow ---"

# Step 1: Login to WandB using the secret
# Make sure you have your WANDB_API_KEY set in the Space secrets
wandb login $WANDB_API_KEY
echo "Step 1 complete: Logged into WandB."

# Step 2: Run the data curation and processing scripts
# (Assuming your curation and processing scripts are in the repo)
# python dataset_curation.py
# python dataset/build_abstract_dataset.py
echo "Step 2 complete: Data processing finished."

# Step 3: Launch the training job
echo "--- LAUNCHING TRAINING ---"
OMP_NUM_THREADS=8 python pretrain.py \
  data_path=data/abstract_optimizer_processed \
  arch=hrm_abstract_optimizer \
  epochs=20000 \
  global_batch_size=32 \
  eval_interval=1000

echo "--- TRAINING COMPLETE ---"

# Step 4: Keep the space running after the script finishes (optional)
# echo "Workflow finished. Idling..."
# sleep infinity