ComputerAssignments/CA8_neural_program_synthesis/scripts/main.py

#!/usr/bin/env python3

import argparse
import torch
import numpy as np
import json
import os
import sys
from typing import Dict, List, Tuple
from pathlib import Path


from neural_program_synthesis import (
    Seq2SeqProgramSynthesizer,
    TransformerProgramSynthesizer,
    CodeGraphEncoder,
    MultiModalCodeEncoder,
    TemplateSynthesizer,
    Seq2SeqTrainer,
    ProgramEvaluator,
    PythonExecutor,
    JavaScriptExecutor,
    SimpleTokenizer,
    SimpleTemplate,
    compare_programs,
)


from train import train_model, save_checkpoint, load_checkpoint
from eval import evaluate_model, generate_programs
from data_generator import generate_synthetic_data, create_dataset

def setup_args() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Neural Program Synthesis Pipeline",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )


    parser.add_argument(
        "--model_type",
        type=str,
        default="seq2seq",
        choices=["seq2seq", "transformer", "graph"],
        help="Type of neural architecture to use",
    )
    parser.add_argument(
        "--vocab_size", type=int, default=1000, help="Vocabulary size for tokenization"
    )
    parser.add_argument(
        "--embedding_dim", type=int, default=256, help="Embedding dimension"
    )
    parser.add_argument(
        "--hidden_dim", type=int, default=512, help="Hidden dimension for RNN/LSTM"
    )
    parser.add_argument(
        "--num_layers", type=int, default=2, help="Number of layers in model"
    )
    parser.add_argument(
        "--nhead",
        type=int,
        default=8,
        help="Number of attention heads (for transformer)",
    )


    parser.add_argument(
        "--batch_size", type=int, default=32, help="Batch size for training"
    )
    parser.add_argument(
        "--learning_rate", type=float, default=0.001, help="Learning rate"
    )
    parser.add_argument(
        "--num_epochs", type=int, default=100, help="Number of training epochs"
    )
    parser.add_argument(
        "--max_length", type=int, default=100, help="Maximum sequence length"
    )


    parser.add_argument(
        "--train_size", type=int, default=1000, help="Number of training samples"
    )
    parser.add_argument(
        "--val_size", type=int, default=200, help="Number of validation samples"
    )
    parser.add_argument(
        "--test_size", type=int, default=300, help="Number of test samples"
    )


    parser.add_argument(
        "--mode",
        type=str,
        default="full",
        choices=["train", "eval", "generate", "full"],
        help="Execution mode",
    )
    parser.add_argument(
        "--checkpoint", type=str, default=None, help="Path to model checkpoint"
    )
    parser.add_argument(
        "--output_dir", type=str, default="./results", help="Directory to save results"
    )
    parser.add_argument(
        "--device", type=str, default="auto", help="Device to use (auto/cpu/cuda)"
    )


    parser.add_argument(
        "--task",
        type=str,
        default="arithmetic",
        choices=["arithmetic", "string", "list_ops", "custom"],
        help="Type of program synthesis task",
    )
    parser.add_argument(
        "--specification",
        type=str,
        default=None,
        help="Natural language specification for generation",
    )

    return parser

def setup_device(device_str: str) -> torch.device:
    if device_str == "auto":
        if torch.cuda.is_available():
            device = torch.device("cuda")
            print("Using CUDA device")
        else:
            device = torch.device("cpu")
            print("Using CPU device")
    else:
        device = torch.device(device_str)
        print(f"Using device: {device}")

    return device

def setup_model(args) -> Tuple[torch.nn.Module, torch.optim.Optimizer, torch.device]:
    device = setup_device(args.device)


    tokenizer = SimpleTokenizer(args.vocab_size)


    if args.model_type == "seq2seq":
        model = Seq2SeqProgramSynthesizer(
            vocab_size=args.vocab_size,
            embedding_dim=args.embedding_dim,
            hidden_dim=args.hidden_dim,
            num_layers=args.num_layers,
        )
    elif args.model_type == "transformer":
        model = TransformerProgramSynthesizer(
            vocab_size=args.vocab_size,
            d_model=args.embedding_dim,
            nhead=args.nhead,
            num_layers=args.num_layers,
        )
    elif args.model_type == "graph":
        model = CodeGraphEncoder(
            vocab_size=args.vocab_size,
            embedding_dim=args.embedding_dim,
            hidden_dim=args.hidden_dim,
        )
    else:
        raise ValueError(f"Unsupported model type: {args.model_type}")

    model = model.to(device)


    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)


    if args.checkpoint and os.path.exists(args.checkpoint):
        checkpoint = load_checkpoint(args.checkpoint)
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        print(f"Loaded checkpoint from {args.checkpoint}")

    return model, optimizer, device, tokenizer

def create_arithmetic_templates() -> List[SimpleTemplate]:
    templates = [
        SimpleTemplate("result = {a} + {b}", ["a", "b"]),
        SimpleTemplate("result = {a} - {b}", ["a", "b"]),
        SimpleTemplate("result = {a} * {b}", ["a", "b"]),
        SimpleTemplate("result = {a} / {b}", ["a", "b"]),
        SimpleTemplate("result = pow({a} * {b}, 2)", ["a", "b"]),
        SimpleTemplate("result = abs({a} - {b})", ["a", "b"]),
    ]
    return templates

def setup_arithmetic_task(args) -> Dict:
    print("Setting up arithmetic program synthesis task...")


    templates = create_arithmetic_templates()


    train_data = generate_synthetic_data(
        task="arithmetic",
        num_samples=args.train_size,
        templates=templates,
        tokenizer_vocab_size=args.vocab_size,
    )

    val_data = generate_synthetic_data(
        task="arithmetic",
        num_samples=args.val_size,
        templates=templates,
        tokenizer_vocab_size=args.vocab_size,
    )

    test_data = generate_synthetic_data(
        task="arithmetic",
        num_samples=args.test_size,
        templates=templates,
        tokenizer_vocab_size=args.vocab_size,
    )


    train_dataset = create_dataset(train_data, args.max_length)
    val_dataset = create_dataset(val_data, args.max_length)
    test_dataset = create_dataset(test_data, args.max_length)

    return {
        "train_dataset": train_dataset,
        "val_dataset": val_dataset,
        "test_dataset": test_dataset,
        "templates": templates,
    }

def run_training_pipeline(args, model, optimizer, device, tokenizer, task_data):
    print("Starting training pipeline...")


    if args.model_type in ["seq2seq", "transformer"]:
        trainer = Seq2SeqTrainer(model, optimizer, device)
    else:

        trainer = None
        print(f"Warning: No specific trainer for {args.model_type}, using generic")


    train_history = train_model(
        model=model,
        trainer=trainer,
        train_dataset=task_data["train_dataset"],
        val_dataset=task_data["val_dataset"],
        epochs=args.num_epochs,
        batch_size=args.batch_size,
        device=device,
        save_dir=args.output_dir,
    )

    print("Training completed!")
    return train_history

def run_evaluation_pipeline(args, model, device, tokenizer, task_data):
    print("Starting evaluation pipeline...")


    evaluator = ProgramEvaluator()


    eval_results = evaluate_model(
        model=model,
        test_dataset=task_data["test_dataset"],
        evaluator=evaluator,
        device=device,
        batch_size=args.batch_size,
    )

    print("Evaluation completed!")
    return eval_results

def run_generation_pipeline(args, model, device, tokenizer):
    print("Starting program generation pipeline...")

    if args.specification:
        print(f"Generating program for: {args.specification}")


        spec_tokens = tokenizer.tokenize(args.specification)
        spec_tensor = torch.tensor([spec_tokens], dtype=torch.long).to(device)
        spec_lengths = torch.tensor([len(spec_tokens)], dtype=torch.long)


        generated_programs = generate_programs(
            model=model,
            input_specifications=spec_tensor,
            input_lengths=spec_lengths,
            tokenizer=tokenizer,
            device=device,
            max_length=args.max_length,
        )

        print("Generated programs:")
        for i, program in enumerate(generated_programs):
            print(f"Program {i+1}: {program}")

        return generated_programs
    else:
        print("No specification provided. Please provide --specification")
        return []

def save_results(args, results: Dict):
    os.makedirs(args.output_dir, exist_ok=True)


    results_file = os.path.join(args.output_dir, "results.json")
    with open(results_file, "w") as f:
        json.dump(results, f, indent=2)


    config = vars(args)
    config_file = os.path.join(args.output_dir, "config.json")
    with open(config_file, "w") as f:
        json.dump(config, f, indent=2)

    print(f"Results saved to {args.output_dir}")

def main():
    parser = setup_args()
    args = parser.parse_args()

    print("=" * 60)
    print("NEURAL PROGRAM SYNTHESIS PIPELINE")
    print("=" * 60)
    print(f"Task: {args.task}")
    print(f"Model: {args.model_type}")
    print(f"Mode: {args.mode}")
    print("=" * 60)

    try:

        model, optimizer, device, tokenizer = setup_model(args)


        task_data = setup_arithmetic_task(args)

        results = {}

        if args.mode in ["train", "full"]:

            train_history = run_training_pipeline(
                args, model, optimizer, device, tokenizer, task_data
            )
            results["training"] = train_history

        if args.mode in ["eval", "full"]:

            eval_results = run_evaluation_pipeline(
                args, model, device, tokenizer, task_data
            )
            results["evaluation"] = eval_results

        if args.mode in ["generate", "full"]:

            generated_programs = run_generation_pipeline(args, model, device, tokenizer)
            results["generation"] = generated_programs


        save_results(args, results)

        print("\n" + "=" * 60)
        print("PIPELINE COMPLETED SUCCESSFULLY!")
        print("=" * 60)

    except Exception as e:
        print(f"\nError during execution: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()