#!/bin/bash

# Get current directory
CURRENT_DIR=$(pwd)

# Create merged_models directory if it doesn't exist
mkdir -p "${CURRENT_DIR}/merged_models"

# Find all checkpoint directories
for checkpoint_dir in checkpoint-*; do
    # Skip if not a directory
    if [ ! -d "$checkpoint_dir" ]; then
        continue
    fi
    
    # Extract checkpoint name
    checkpoint_name=$(basename "$checkpoint_dir")
    
    # Create output directory
    mkdir -p "${CURRENT_DIR}/merged_models/${checkpoint_name}"
    
    echo "Processing ${checkpoint_name}..."
    
    # Run the merge command
    python3 -m axolotl.cli.merge_lora model.yml \
        --lora_model_dir="./${checkpoint_name}" \
        --output_dir="${CURRENT_DIR}/merged_models/${checkpoint_name}"
    
    # Check if merge was successful
    if [ -d "${CURRENT_DIR}/merged_models/${checkpoint_name}/merged" ]; then
        # Move all files from merged directory up one level
        mv "${CURRENT_DIR}/merged_models/${checkpoint_name}/merged"/* \
           "${CURRENT_DIR}/merged_models/${checkpoint_name}/"
        
        # Remove empty merged directory
        rmdir "${CURRENT_DIR}/merged_models/${checkpoint_name}/merged"
        
        echo "Successfully processed ${checkpoint_name}"
    else
        echo "Error processing ${checkpoint_name}"
    fi
done

echo "All checkpoints processed!"