ComputerAssignments/CA8_neural_program_synthesis/scripts/eval.py

#!/usr/bin/env python3

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
import json
import os
import time
from typing import Dict, List, Tuple, Optional, Union
from tqdm import tqdm

from neural_program_synthesis import (
    ProgramEvaluator,
    PythonExecutor,
    JavaScriptExecutor,
    compare_programs,
)
from train import create_data_loader

class SynthesisEvaluator:

    def __init__(self, executor_name: str = "python", syntax_check: bool = True):
        self.syntax_check = syntax_check


        if executor_name.lower() == "python":
            self.executor = PythonExecutor()
        elif executor_name.lower() == "javascript":
            self.executor = JavaScriptExecutor()
        else:
            raise ValueError(f"Unsupported executor: {executor_name}")


        self.metrics = {
            "total_samples": 0,
            "syntax_valid": 0,
            "functionally_correct": 0,
            "runtime_errors": 0,
            "execution_success_rate": 0.0,
            "functional_correct_rate": 0.0,
            "overall_success_rate": 0.0,
        }

    def evaluate_syntax(self, generated_program: str) -> bool:
        if not self.syntax_check:
            return True

        try:
            if isinstance(self.executor, PythonExecutor):
                import ast

                ast.parse(generated_program)
                return True
            else:

                return len(generated_program.strip()) > 0
        except Exception:
            return False

    def evaluate_functional(
        self, generated_program: str = str, test_case: Dict = None
    ) -> Tuple[bool, str]:
        if test_case is None:
            test_case = {"input": {}, "expected_output": {}, "test_type": "value"}

        try:

            success, output_data, error_msg = self.executor.execute(
                generated_program, test_case.get("input", {})
            )

            if not success:
                return False, f"Runtime error: {error_msg}"


            expected_output = test_case.get("expected_output", {})
            test_type = test_case.get("test_type", "value")

            if test_type == "boolean":
                match = output_data.get("result", False) == expected_output.get(
                    "result", False
                )
            else:
                match = str(output_data).strip() == str(expected_output).strip()

            return match, "" if match else "Output mismatch"

        except Exception as e:
            return False, f"Evaluation error: {str(e)}"

    def evaluate_batch(
        self, generated_programs: List[str], test_cases: Optional[List[Dict]] = None
    ) -> Dict[str, Union[int, float]]:
        if test_cases is None:
            test_cases = [
                {"input": {}, "expected_output": {}, "test_type": "value"}
                for _ in generated_programs
            ]

        batch_metrics = {
            "total_samples": len(generated_programs),
            "syntax_valid": 0,
            "functionally_correct": 0,
            "runtime_errors": 0,
            "execution_errors": 0,
            "output_mismatches": 0,
        }

        for program, test_case in zip(generated_programs, test_cases):

            syntax_valid = self.evaluate_syntax(program)
            if syntax_valid:
                batch_metrics["syntax_valid"] += 1
            else:
                continue


            func_success, error_msg = self.evaluate_functional(program, test_case)

            if "Runtime error" in error_msg:
                batch_metrics["runtime_errors"] += 1
            elif "Evaluation error" in error_msg:
                batch_metrics["execution_errors"] += 1
            elif "Output mismatch" in error_msg:
                batch_metrics["output_mismatches"] += 1
            elif func_success:
                batch_metrics["functionally_correct"] += 1


        total = batch_metrics["total_samples"]
        batch_metrics["syntax_valid_rate"] = (
            batch_metrics["syntax_valid"] / total if total > 0 else 0.0
        )
        batch_metrics["functional_correct_rate"] = (
            batch_metrics["functionally_correct"] / total if total > 0 else 0.0
        )
        batch_metrics["execution_success_rate"] = (
            batch_metrics["syntax_valid"] / total if total > 0 else 0.0
        )

        return batch_metrics

def generate_programs(
    model: nn.Module,
    input_specifications: torch.Tensor,
    input_lengths: torch.Tensor,
    tokenizer,
    device: torch.device = torch.device("cpu"),
    max_length: int = 100,
    num_candidates: int = 1,
    temperature: float = 1.0,
) -> List[str]:
    model.eval()
    generated_programs = []

    with torch.no_grad():
        for i in range(num_candidates):

            if hasattr(model, "generate"):
                generated_tokens = model.generate(
                    input_specifications,
                    input_lengths,
                    max_length=max_length,
                    bos_id=1,
                    eos_id=2,
                )


                program_tokens = generated_tokens[0].cpu().tolist()

                program_tokens = [
                    token for token in program_tokens if token not in [0, 1, 2]
                ]

                if program_tokens:
                    program_code = tokenizer.detokenize(program_tokens)
                else:
                    program_code = "# No valid program generated"

                generated_programs.append(program_code)


            elif hasattr(model, "transformer"):

                input_embedded = model.embedding(input_specifications)
                generated = input_embedded

                for _ in range(max_length):
                    output = model.transformer(
                        generated.transpose(0, 1), generated.transpose(0, 1)
                    )
                    logits = model.output_proj(output.transpose(0, 1))
                    next_token = torch.argmax(logits[:, -1, :], dim=-1)

                    if next_token.item() == 2:
                        break

                    next_embedded = model.embedding(next_token.unsqueeze(0))
                    generated = torch.cat([generated, next_embedded], dim=1)


                tokens = torch.argmax(output, dim=-1)[0].cpu().tolist()
                program_tokens = [token for token in tokens if token not in [0, 1, 2]]
                program_code = (
                    tokenizer.detokenize(program_tokens)
                    if program_tokens
                    else "# Generation failed"
                )
                generated_programs.append(program_code)

            else:

                generated_programs.append("# Unsupported model type for generation")

    return generated_programs

def evaluate_model(
    model: nn.Module,
    test_dataset,
    evaluator: SynthesisEvaluator,
    device: torch.device = torch.device("cpu"),
    batch_size: int = 32,
    generate_per_sample: int = 1,
) -> Dict[str, Union[int, float]]:
    print(f"Evaluating model on {len(test_dataset)} test samples...")


    test_loader = create_data_loader(test_dataset, batch_size, shuffle=False)

    all_metrics = {
        "total_samples": 0,
        "syntax_valid": 0,
        "functionally_correct": 0,
        "runtime_errors": 0,
        "execution_errors": 0,
        "output_mismatches": 0,
    }

    generated_programs_log = []

    model.eval()
    with torch.no_grad():
        for batch_idx, batch in enumerate(tqdm(test_loader, desc="Evaluating")):

            batch = {
                k: v.to(device) if isinstance(v, torch.Tensor) else v
                for k, v in batch.items()
            }

            batch_size_actual = batch["input_ids"].size(0)


            batch_programs = []
            batch_test_cases = []

            for i in range(batch_size_actual):

                sample_programs = []
                for _ in range(generate_per_sample):
                    input_ids = batch["input_ids"][i : i + 1]
                    input_length = batch["input_lengths"][i : i + 1]


                    if hasattr(model, "generate"):
                        generated_tokens = model.generate(
                            input_ids, input_length, max_length=50
                        )
                    else:
                        generated_tokens = torch.randint(3, 100, (1, 10))


                    tokens = generated_tokens[0].cpu().tolist()
                    tokens = [t for t in tokens if t not in [0, 1, 2]]
                    program_code = f"# Generated program {i}_{len(sample_programs)}\nprint('Result:', {tokens[0] if tokens else 42})"
                    sample_programs.append(program_code)

                batch_programs.extend(sample_programs)


                test_case = {
                    "input": {},
                    "expected_output": {"result": 42},
                    "test_type": "value",
                }
                batch_test_cases.extend([test_case] * generate_per_sample)


            batch_metrics = evaluator.evaluate_batch(batch_programs, batch_test_cases)


            for key in all_metrics.keys():
                if key in batch_metrics:
                    all_metrics[key] += batch_metrics[key]


            generated_programs_log.extend(batch_programs)


    total = all_metrics["total_samples"]
    if total > 0:
        all_metrics["syntax_valid_rate"] = all_metrics["syntax_valid"] / total
        all_metrics["functional_correct_rate"] = (
            all_metrics["functionally_correct"] / total
        )
        all_metrics["execution_success_rate"] = all_metrics["syntax_valid"] / total
        all_metrics["overall_success_rate"] = (
            all_metrics["functionally_correct"] / total
        ) * (all_metrics["syntax_valid"] / total)
    else:
        all_metrics.update(
            {
                "syntax_valid_rate": 0.0,
                "functional_correct_rate": 0.0,
                "execution_success_rate": 0.0,
                "overall_success_rate": 0.0,
            }
        )


    print(f"\n" + "=" * 60)
    print("EVALUATION SUMMARY")
    print("=" * 60)
    print(f"Total samples: {all_metrics['total_samples']}")
    print(
        f"Syntax valid: {all_metrics['syntax_valid']} ({all_metrics['syntax_valid_rate']:.2%})"
    )
    print(
        f"Functionally correct: {all_metrics['functionally_correct']} ({all_metrics['functional_correct_rate']:.2%})"
    )
    print(f"Runtime errors: {all_metrics['runtime_errors']}")
    print(f"Execution errors: {all_metrics['execution_errors']}")
    print(f"Output mismatches: {all_metrics['output_mismatches']}")
    print(f"Overall success rate: {all_metrics['overall_success_rate']:.2%}")
    print("=" * 60)


    all_metrics["generated_programs"] = generated_programs_log[
        :10
    ]

    return all_metrics

def benchmark_model(
    model: nn.Module,
    test_dataset,
    tokenizer,
    device: torch.device = torch.device("cpu"),
    num_runs: int = 3,
) -> Dict[str, Union[float, List[float]]]:
    print(f"Benchmarking model across {num_runs} runs...")

    evaluator = SynthesisEvaluator()

    all_metrics = []
    execution_times = []

    for run in range(num_runs):
        print(f"Benchmark run {run + 1}/{num_runs}")

        start_time = time.time()


        metrics = evaluate_model(model, test_dataset, evaluator, device)

        execution_time = time.time() - start_time

        all_metrics.append(metrics)
        execution_times.append(execution_time)

        print(f"Run {run + 1} completed in {execution_time:.2f}s")


    aggregated = {
        "runs": num_runs,
        "execution_times": execution_times,
        "avg_execution_time": np.mean(execution_times),
        "std_execution_time": np.std(execution_times),
    }


    for metric in [
        "syntax_valid_rate",
        "functional_correct_rate",
        "execution_success_rate",
        "overall_success_rate",
    ]:
        values = [m[metric] for m in all_metrics]
        aggregated[f"avg_{metric}"] = np.mean(values)
        aggregated[f"std_{metric}"] = np.std(values)

    print(f"\nBenchmark Summary:")
    print(
        f"Average execution time: {aggregated['avg_execution_time']:.2f}s ± {aggregated['std_execution_time']:.2f}s"
    )
    print(
        f"Average overall success rate: {aggregated['avg_overall_success_rate']:.2%} ± {aggregated['std_overall_success_rate']:.2%}"
    )

    return aggregated

def save_evaluation_results(results: Dict, output_path: str):
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    with open(output_path, "w") as f:
        json.dump(results, f, indent=2, default=str)

    print(f"Evaluation results saved to {output_path}")

if __name__ == "__main__":
    print("Evaluation module for Neural Program Synthesis")
    print("This module should be imported and used by main.py")