ComputerAssignments/CA8_neural_program_synthesis/scripts/run.sh

#!/bin/bash



set -e


SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PYTHON="python3"
PIP="pip3"


RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m'

echo -e "${BLUE}================================================${NC}"
echo -e "${BLUE}    NEURAL PROGRAM SYNTHESIS PIPELINE${NC}"
echo -e "${BLUE}    Complete Implementation${NC}"
echo -e "${BLUE}================================================${NC}"


print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_step() {
    echo -e "${PURPLE}[STEP]${NC} $1"
}

print_result() {
    echo -e "${CYAN}[RESULT]${NC} $1"
}


check_python() {
    print_info "Checking Python installation..."
    if ! command -v $PYTHON &> /dev/null; then
        print_error "Python 3 is required but not installed. Please install Python 3."
        exit 1
    fi
    
    PYTHON_VERSION=$($PYTHON --version 2>&1 | cut -d' ' -f2 | cut -d'.' -f1,2)
    print_success "Python $PYTHON_VERSION found"
}


check_pytorch() {
    print_info "Checking PyTorch installation..."
    if $PYTHON -c "import torch; print(f'PyTorch {torch.__version__} found')" 2>/dev/null; then
        print_success "PyTorch is available"
    else
        print_warning "PyTorch not found. Installing..."
        install_dependencies
    fi
}


install_dependencies() {
    print_info "Installing dependencies..."
    if [ ! -f "requirements.txt" ]; then
        print_error "requirements.txt not found!"
        exit 1
    fi
    
    $PIP install -r requirements.txt
    print_success "Dependencies installed successfully"
}


create_directories() {
    print_info "Creating necessary directories..."
    mkdir -p results
    mkdir -p visualizations
    mkdir -p checkpoints
    mkdir -p logs
    mkdir -p data
    print_success "Directories created"
}


generate_data() {
    print_step "Generating synthetic training data..."
    
    local data_args=""
    if [ ! -z "$TASK" ]; then
        data_args="$data_args --task $TASK"
    fi
    
    if [ ! -z "$TRAIN_SIZE" ]; then
        data_args="$data_args --train_size $TRAIN_SIZE"
    fi
    
    if [ ! -z "$VAL_SIZE" ]; then
        data_args="$data_args --val_size $VAL_SIZE"
    fi
    
    if [ ! -z "$TEST_SIZE" ]; then
        data_args="$data_args --test_size $TEST_SIZE"
    fi
    
    $PYTHON -c "
from data_generator import generate_synthetic_data, create_splits, save_data
import json

print('Generating synthetic data...')
data = generate_synthetic_data(
    task='$TASK',
    num_samples=$((TRAIN_SIZE + VAL_SIZE + TEST_SIZE)),
    seed=42
)

train_data, val_data, test_data = create_splits(data, train_ratio=0.7, val_ratio=0.15, test_ratio=0.15)

save_data(train_data, 'data/train_data.json')
save_data(val_data, 'data/val_data.json')
save_data(test_data, 'data/test_data.json')

print(f'Generated {len(train_data)} training samples')
print(f'Generated {len(val_data)} validation samples')
print(f'Generated {len(test_data)} test samples')
"
    
    print_success "Data generation completed"
}


run_training() {
    print_step "Starting training pipeline..."
    
    local train_args=""
    if [ ! -z "$TASK" ]; then
        train_args="$train_args --task $TASK"
    fi
    
    if [ ! -z "$MODEL_TYPE" ]; then
        train_args="$train_args --model_type $MODEL_TYPE"
    fi
    
    if [ ! -z "$EPOCHS" ]; then
        train_args="$train_args --num_epochs $EPOCHS"
    fi
    
    if [ ! -z "$BATCH_SIZE" ]; then
        train_args="$train_args --batch_size $BATCH_SIZE"
    fi
    
    if [ ! -z "$LEARNING_RATE" ]; then
        train_args="$train_args --learning_rate $LEARNING_RATE"
    fi
    
    train_args="$train_args --output_dir ./results"
    
    print_info "Training command: $PYTHON main.py --mode train $train_args"
    $PYTHON main.py --mode train $train_args 2>&1 | tee logs/training.log
    
    if [ $? -eq 0 ]; then
        print_success "Training completed successfully"
    else
        print_error "Training failed. Check logs/training.log for details."
        exit 1
    fi
}


run_evaluation() {
    print_step "Starting evaluation pipeline..."
    
    local eval_args=""
    if [ ! -z "$CHECKPOINT" ]; then
        eval_args="$eval_args --checkpoint $CHECKPOINT"
    else

        if [ -f "results/best_model.pt" ]; then
            CHECKPOINT="results/best_model.pt"
        elif [ -f "results/final_model.pt" ]; then
            CHECKPOINT="results/final_model.pt"
        else
            print_error "No checkpoint found for evaluation"
            exit 1
        fi
    fi
    
    eval_args="$eval_args --checkpoint $CHECKPOINT"
    
    if [ ! -z "$TASK" ]; then
        eval_args="$eval_args --task $TASK"
    fi
    
    eval_args="$eval_args --output_dir ./results"
    
    print_info "Evaluation command: $PYTHON main.py --mode eval $eval_args"
    $PYTHON main.py --mode eval $eval_args 2>&1 | tee logs/evaluation.log
    
    if [ $? -eq 0 ]; then
        print_success "Evaluation completed successfully"
    else
        print_error "Evaluation failed. Check logs/evaluation.log for details."
        exit 1
    fi
}


run_generation() {
    print_step "Starting program generation..."
    
    if [ -z "$SPECIFICATION" ]; then
        SPECIFICATION="add two numbers and return the result"
        print_warning "No specification provided, using default: '$SPECIFICATION'"
    fi
    
    local gen_args=""
    if [ ! -z "$CHECKPOINT" ]; then
        gen_args="$gen_args --checkpoint $CHECKPOINT"
    else
        if [ -f "results/best_model.pt" ]; then
            CHECKPOINT="results/best_model.pt"
        elif [ -f "results/final_model.pt" ]; then
            CHECKPOINT="results/final_model.pt"
        else
            print_error "No checkpoint found for generation"
            exit 1
        fi
    fi
    
    gen_args="$gen_args --checkpoint $CHECKPOINT"
    
    if [ ! -z "$MODEL_TYPE" ]; then
        gen_args="$gen_args --model_type $MODEL_TYPE"
    fi
    
    gen_args="$gen_args --specification \"$SPECIFICATION\""
    gen_args="$gen_args --output_dir ./results"
    
    print_info "Generation command: $PYTHON main.py --mode generate $gen_args"
    $PYTHON main.py --mode generate $gen_args 2>&1 | tee logs/generation.log
    
    if [ $? -eq 0 ]; then
        print_success "Program generation completed successfully"
    else
        print_error "Program generation failed. Check logs/generation.log for details."
        exit 1
    fi
}


create_visualizations() {
    print_step "Creating visualizations..."
    
    $PYTHON -c "
from visualizations import SynthesisVisualizer
import json
import os

visualizer = SynthesisVisualizer('./visualizations')


results = {}
if os.path.exists('results/training_results.json'):
    with open('results/training_results.json', 'r') as f:
        results.update(json.load(f))

if os.path.exists('results/evaluation_results.json'):
    with open('results/evaluation_results.json', 'r') as f:
        results.update(json.load(f))

if os.path.exists('results/generation_results.json'):
    with open('results/generation_results.json', 'r') as f:
        results.update(json.load(f))


if 'train_losses' in results and 'val_losses' in results:
    visualizer.plot_training_curves(
        results['train_losses'],
        results['val_losses'],
        results.get('train_accs'),
        results.get('val_accs')
    )
    print('Training curves visualization created')

if results:

    metrics_to_plot = {k: v for k, v in results.items() 
                      if isinstance(v, (int, float)) and 'rate' in k.lower()}
    if metrics_to_plot:
        visualizer.plot_evaluation_metrics(metrics_to_plot)
        print('Evaluation metrics visualization created')
    

    visualizer.create_summary_report(results)
    visualizer.save_results_json(results)
    print('Summary report created')

print('All visualizations completed successfully!')
"
    
    print_success "Visualizations created"
}


run_full_pipeline() {
    print_step "Starting complete neural program synthesis pipeline..."
    

    create_directories
    

    generate_data
    

    run_training
    

    run_evaluation
    

    run_generation
    

    create_visualizations
    
    print_success "Complete pipeline finished successfully!"
    print_result "Results saved in:"
    print_result "  - Training: ./results/training_results.json"
    print_result "  - Evaluation: ./results/evaluation_results.json"
    print_result "  - Generation: ./results/generation_results.json"
    print_result "  - Visualizations: ./visualizations/"
    print_result "  - Logs: ./logs/"
}


test_components() {
    print_step "Testing individual components..."
    

    print_info "Testing data generator..."
    $PYTHON -c "
from data_generator import generate_synthetic_data
data = generate_synthetic_data(num_samples=10, task='arithmetic')
print(f'Generated {len(data)} samples successfully')
"
    

    print_info "Testing model creation..."
    $PYTHON -c "
from neural_program_synthesis import Seq2SeqProgramSynthesizer
import torch

model = Seq2SeqProgramSynthesizer(vocab_size=1000)
print(f'Model created successfully: {model}')
print(f'Model parameters: {sum(p.numel() for p in model.parameters())}')
"
    

    print_info "Testing evaluator..."
    $PYTHON -c "
from eval import SynthesisEvaluator
evaluator = SynthesisEvaluator()
print('Evaluator created successfully')
"
    
    print_success "All components tested successfully"
}


cleanup() {
    print_info "Cleaning up old results..."
    rm -rf ./results ./visualizations ./checkpoints ./logs ./data
    print_success "Cleanup completed"
}


show_help() {
    echo "Neural Program Synthesis Pipeline Runner"
    echo ""
    echo "Usage: $0 [OPTIONS] [MODE]"
    echo ""
    echo "MODES:"
    echo "  setup       Setup environment and install dependencies"
    echo "  data        Generate synthetic training data"
    echo "  train       Train the neural program synthesis model"
    echo "  eval        Evaluate trained model"
    echo "  generate    Generate programs from natural language"
    echo "  visualize   Create visualizations from results"
    echo "  full        Run complete training + evaluation + generation pipeline"
    echo "  test        Test individual components"
    echo "  cleanup     Remove old result directories"
    echo "  help        Show this help message"
    echo ""
    echo "OPTIONS:"
    echo "  --task ARITHMETIC        Task type (arithmetic, string, list_ops)"
    echo "  --model_type SEQ2SEQ     Model architecture (seq2seq, transformer, graph)"
    echo "  --epochs 100             Number of training epochs"
    echo "  --batch_size 32          Training batch size"
    echo "  --learning_rate 0.001    Learning rate for training"
    echo "  --checkpoint PATH        Path to model checkpoint"
    echo "  --specification TEXT     Natural language specification for generation"
    echo "  --train_size 1000        Number of training samples"
    echo "  --val_size 200           Number of validation samples"
    echo "  --test_size 300          Number of test samples"
    echo "  --vocab_size 1000        Vocabulary size"
    echo ""
    echo "EXAMPLES:"
    echo "  $0 setup"
    echo "  $0 test"
    echo "  $0 data --task arithmetic --train_size 1000"
    echo "  $0 train --task arithmetic --epochs 50 --batch_size 64"
    echo "  $0 eval --checkpoint ./results/best_model.pt"
    echo "  $0 generate --specification \"multiply two numbers\""
    echo "  $0 full --task arithmetic --epochs 100 --batch_size 32"
    echo "  $0 visualize"
}


main() {
    cd "$SCRIPT_DIR"
    
    MODE=""
    TASK="arithmetic"
    MODEL_TYPE="seq2seq"
    EPOCHS="50"
    BATCH_SIZE="32"
    LEARNING_RATE="0.001"
    CHECKPOINT=""
    SPECIFICATION=""
    TRAIN_SIZE="1000"
    VAL_SIZE="200"
    TEST_SIZE="300"
    VOCAB_SIZE="1000"
    

    while [[ $
        case $1 in
            setup|data|train|eval|generate|visualize|full|test|cleanup|help)
                MODE="$1"
                shift
                ;;
            --task)
                TASK="$2"
                shift 2
                ;;
            --model_type)
                MODEL_TYPE="$2"
                shift 2
                ;;
            --epochs)
                EPOCHS="$2"
                shift 2
                ;;
            --batch_size)
                BATCH_SIZE="$2"
                shift 2
                ;;
            --learning_rate)
                LEARNING_RATE="$2"
                shift 2
                ;;
            --checkpoint)
                CHECKPOINT="$2"
                shift 2
                ;;
            --specification)
                SPECIFICATION="$2"
                shift 2
                ;;
            --train_size)
                TRAIN_SIZE="$2"
                shift 2
                ;;
            --val_size)
                VAL_SIZE="$2"
                shift 2
                ;;
            --test_size)
                TEST_SIZE="$2"
                shift 2
                ;;
            --vocab_size)
                VOCAB_SIZE="$2"
                shift 2
                ;;
            *)
                print_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    if [ -z "$MODE" ]; then
        print_error "Please specify a mode"
        show_help
        exit 1
    fi
    
    case $MODE in
        setup)
            check_python
            check_pytorch
            install_dependencies
            create_directories
            print_success "Setup completed successfully!"
            ;;
        data)
            check_python
            create_directories
            generate_data
            ;;
        train)
            check_python
            check_pytorch
            run_training
            ;;
        eval)
            check_python
            check_pytorch
            run_evaluation
            ;;
        generate)
            check_python
            check_pytorch
            run_generation
            ;;
        visualize)
            check_python
            create_visualizations
            ;;
        full)
            check_python
            check_pytorch
            run_full_pipeline
            ;;
        test)
            check_python
            check_pytorch
            test_components
            ;;
        cleanup)
            cleanup
            ;;
        help)
            show_help
            ;;
        *)
            print_error "Unknown mode: $MODE"
            show_help
            exit 1
            ;;
    esac
}


main "$@"