run_multiple_models.sh

#!/bin/bash

# Script to iterate through multiple LLM models for function generation
# Usage: ./run_multiple_models.sh [additional_args...]

set -e  # Exit on any error

# Define the models to iterate through
MODELS=("gpt-4o" "gpt-4o-mini" "o3" "DeepSeek-V3-Online" "Qwen2.5-7B-Instruct" "Qwen2.5-72B-Instruct")
# MODELS=("Qwen2.5-7B-Instruct" "o3" "DeepSeek-V3-Online" "Qwen2.5-72B-Instruct")

# MODELS=("Qwen2.5-7B-Instruct" "Qwen2.5-72B-Instruct")

# Default arguments (can be overridden by command line)
DEFAULT_INPUT_FILE="Jsons/extracted_functions_with_comments.json"
DEFAULT_CONFIG_FILE="openai_config.json"
DEFAULT_BATCH_SIZE=1
DEFAULT_MAX_FUNCTIONS=""

# Base output directory for results
OUTPUT_DIR="results_snake"
mkdir -p "$OUTPUT_DIR"

# Colors for output formatting
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Function to print colored output
print_status() {
    echo -e "${BLUE}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[$(date '+%Y-%m-%d %H:%M:%S')] ✅${NC} $1"
}

print_error() {
    echo -e "${RED}[$(date '+%Y-%m-%d %H:%M:%S')] ❌${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[$(date '+%Y-%m-%d %H:%M:%S')] ⚠️${NC} $1"
}

# Function to check if Python script exists
check_script() {
    if [ ! -f "generate_functions_from_descriptions.py" ]; then
        print_error "generate_functions_from_descriptions.py not found in current directory!"
        exit 1
    fi
}

# Function to run generation for a specific model
run_model_generation() {
    local model=$1
    local timestamp=$(date '+%Y%m%d_%H%M%S')
    local output_file="${OUTPUT_DIR}/patches_${model//[-.]/_}_${timestamp}.json"
    
    print_status "Starting function generation with model: ${model}"
    print_status "Output will be saved to: ${output_file}"
    
    # Build the command
    local cmd="python3 generate_functions_from_descriptions.py"
    cmd="$cmd --model '$model'"
    cmd="$cmd --output-file '$output_file'"
    
    # Add additional arguments passed to this script
    if [ $# -gt 1 ]; then
        shift  # Remove the model argument
        cmd="$cmd $*"
    fi
    
    # Add default arguments if not already specified
    if [[ "$*" != *"--input-file"* ]]; then
        cmd="$cmd --input-file '$DEFAULT_INPUT_FILE'"
    fi
    
    if [[ "$*" != *"--config-file"* ]]; then
        cmd="$cmd --config-file '$DEFAULT_CONFIG_FILE'"
    fi
    
    if [[ "$*" != *"--batch-size"* ]]; then
        cmd="$cmd --batch-size $DEFAULT_BATCH_SIZE"
    fi
    
    if [[ "$*" != *"--max-functions"* && -n "$DEFAULT_MAX_FUNCTIONS" ]]; then
        cmd="$cmd --max-functions $DEFAULT_MAX_FUNCTIONS"
    fi
    
    print_status "Executing: $cmd"
    
    # Execute the command and capture output
    if eval "$cmd" 2>&1 | tee "${OUTPUT_DIR}/log_${model//[-.]/_}_${timestamp}.txt"; then
        print_success "Successfully completed generation with model: ${model}"
        print_success "Results saved to: ${output_file}"
        return 0
    else
        print_error "Failed to complete generation with model: ${model}"
        return 1
    fi
}

# Main execution
main() {
    local start_time=$(date +%s)
    local successful_models=()
    local failed_models=()
    
    print_status "🚀 Starting multi-model function generation"
    print_status "Models to process: ${MODELS[*]}"
    print_status "Results will be saved to: $OUTPUT_DIR/"
    
    # Check prerequisites
    check_script
    
    # Verify config file exists
    if [ ! -f "$DEFAULT_CONFIG_FILE" ]; then
        print_warning "Config file $DEFAULT_CONFIG_FILE not found. The script will create a template."
    fi
    
    echo -e "\n" "="*80
    echo "MULTI-MODEL FUNCTION GENERATION"
    echo "="*80
    
    # Iterate through each model
    for model in "${MODELS[@]}"; do
        echo -e "\n" "-"*60
        print_status "Processing model: $model"
        echo "-"*60
        
        if run_model_generation "$model" "$@"; then
            successful_models+=("$model")
        else
            failed_models+=("$model")
            print_warning "Continuing with next model..."
        fi
        
        # # Add delay between models to avoid rate limiting
        # if [ "$model" != "${MODELS[-1]}" ]; then
        #     print_status "Waiting 10 seconds before next model..."
        #     sleep 10
        # fi
    done
    
    # Calculate execution time
    local end_time=$(date +%s)
    local execution_time=$((end_time - start_time))
    local minutes=$((execution_time / 60))
    local seconds=$((execution_time % 60))
    
    # Print final summary
    echo -e "\n" "="*80
    echo "EXECUTION SUMMARY"
    echo "="*80
    
    print_status "Total execution time: ${minutes}m ${seconds}s"
    print_status "Successful models: ${#successful_models[@]}/${#MODELS[@]}"
    
    if [ ${#successful_models[@]} -gt 0 ]; then
        print_success "Successfully processed models:"
        for model in "${successful_models[@]}"; do
            echo "  ✅ $model"
        done
    fi
    
    if [ ${#failed_models[@]} -gt 0 ]; then
        print_error "Failed models:"
        for model in "${failed_models[@]}"; do
            echo "  ❌ $model"
        done
    fi
    
    print_status "Results directory: $OUTPUT_DIR/"
    print_status "Check individual log files for detailed information"
    
    # List generated files
    echo -e "\n" "📁 Generated files:"
    ls -la "$OUTPUT_DIR/"*.json 2>/dev/null || echo "No JSON files generated"
    
    # Return appropriate exit code
    if [ ${#failed_models[@]} -eq 0 ]; then
        print_success "🎉 All models completed successfully!"
        exit 0
    elif [ ${#successful_models[@]} -gt 0 ]; then
        print_warning "⚠️ Some models completed successfully, but some failed"
        exit 2
    else
        print_error "💥 All models failed!"
        exit 1
    fi
}

# Help function
show_help() {
    cat << EOF
Usage: $0 [OPTIONS]

This script iterates through multiple LLM models (gpt-4o, gpt-4o-mini, o3) 
and generates functions using the generate_functions_from_descriptions.py script.

Options:
  --input-file FILE         Input JSON file with function descriptions
                           (default: $DEFAULT_INPUT_FILE)
  --config-file FILE        OpenAI configuration file
                           (default: $DEFAULT_CONFIG_FILE)
  --batch-size N           Number of functions to process in each batch
                           (default: $DEFAULT_BATCH_SIZE)
  --max-functions N        Maximum number of functions to process
  --repo-path PATH         Path to repository for enhanced context
  --use-scot               Use Structured Chain-of-Thought prompting
  --no-scot                Disable SCoT prompting
  --help                   Show this help message

Examples:
  $0                                    # Run with defaults
  $0 --max-functions 10                 # Limit to 10 functions per model
  $0 --repo-path /path/to/repo          # Enable repository context
  $0 --batch-size 5 --use-scot          # Custom batch size with SCoT

Results are saved to the 'results/' directory with timestamps.
Each model gets its own output file and log file.

EOF
}

# Check for help flag
if [[ "$1" == "--help" || "$1" == "-h" ]]; then
    show_help
    exit 0
fi

# Run main function with all arguments
main "$@"