evaluate_post_trained.py

from transformers import AutoModelForCausalLM, AutoTokenizer
from StarCoder import generate_code, execute
from datasets import load_dataset
import json
import os
import re
import matplotlib.pyplot as plt
import numpy as np
import argparse

def evaluate_post_trained_model(percentage=1, num_problems=164):
    """Evaluate an already post-trained model on HumanEval"""
    print(f"\n=== Evaluating Post-Trained Model ({percentage}%) ===")
    
    # Load existing model
    model_path = f'results/post_trained_{percentage}pct'
    if not os.path.exists(model_path):
        print(f"Error: Model not found at {model_path}")
        return None
        
    print(f"Loading model from {model_path}...")
    model = AutoModelForCausalLM.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # Load HumanEval dataset
    humaneval = load_dataset("openai_humaneval")
    problems = humaneval['test'].select(range(num_problems))
    
    results = []
    
    # Improve prompt engineering for better results
    for i, problem in enumerate(problems):
        print(f"\nProblem {i+1}/{num_problems}: {problem['task_id']}")
        
        # Better prompt that constrains output
        enhanced_prompt = f"""
            Complete the following Python function. Provide ONLY the implementation, no tests or comments.
            Make sure your solution is complete and ends with a return statement.

            {problem['prompt']}

            Implementation:
            """
        
        # Generate and preprocess code
        generated_code = generate_code(model, tokenizer, enhanced_prompt)
        
        # Add missing imports for type annotations
        if any(typ in generated_code for typ in ["List", "Optional", "Tuple", "Any"]):
            if "from typing import" not in generated_code:
                generated_code = "from typing import List, Optional, Tuple, Any\n\n" + generated_code
        
        # Clean up generated code
        cleaned_code = clean_generated_code(generated_code)
        
        # Execute tests with processed code
        success, output = execute(cleaned_code, problem['test'])
        
        results.append({
            'task_id': problem['task_id'],
            'prompt': problem['prompt'],
            'generated_code': cleaned_code,
            'execution_success': success,
            'output': output,
            'test_cases': problem['test']
        })
        
        print(f"Success: {success}")
    
    # Save evaluation results
    os.makedirs(f"{model_path}", exist_ok=True)
    results_path = f"{model_path}/humaneval_results.json"
    with open(results_path, 'w') as f:
        json.dump(results, f, indent=2)
    
    # Enhanced summary with comparison to previous results
    post_trained_success = sum(1 for r in results if r['execution_success'])
    
    # Load baseline results if available
    baseline_success = 0
    baseline_total = 0
    try:
        with open('results/baseline_results.json', 'r') as f:
            baseline_results = json.load(f)
            baseline_success = sum(1 for r in baseline_results if r['execution_success'])
            baseline_total = len(baseline_results)
    except (FileNotFoundError, json.JSONDecodeError):
        print("No baseline results found for comparison")
    
    # Load self-evaluation results if available
    self_eval_success = 0
    self_eval_total = 0
    try:
        with open('results/self_eval_results.json', 'r') as f:
            self_eval_results = json.load(f)
            # Count improved successes + original successes
            for r in self_eval_results:
                if r.get('improved_success', False) or r.get('execution_success', False):
                    self_eval_success += 1
            self_eval_total = len(self_eval_results)
    except (FileNotFoundError, json.JSONDecodeError):
        print("No self-evaluation results found for comparison")
    
    # Print comprehensive summary
    total = len(results)
    print(f"\n=== RESULTS SUMMARY ===")
    print(f"Baseline model:      {baseline_success}/{baseline_total} ({baseline_success/baseline_total*100:.2f}% success)")
    print(f"Self-evaluated:      {self_eval_success}/{self_eval_total} ({self_eval_success/self_eval_total*100:.2f}% success)")
    print(f"Post-trained ({percentage}%): {post_trained_success}/{total} ({post_trained_success/total*100:.2f}% success)")
    
    if baseline_total > 0:
        improvement = (post_trained_success/total - baseline_success/baseline_total) * 100
        print(f"Improvement over baseline: {improvement:.2f}%")
    
    print(f"\nResults saved to {results_path}")
    
    return results

def clean_generated_code(code):
    """Clean up the generated code for better execution."""
    # Remove notebook artifacts
    code = re.sub(r'\[markdown\].*', '', code)
    code = re.sub(r'# Test.*|# TODO.*', '', code)
    
    # Remove test cases and assertions
    lines = [line for line in code.split('\n') 
             if not line.strip().startswith(('assert', 'sert', 'rt ', 'st '))
             and not line.strip().startswith('>>>')]
    
    # Extract just the function implementation
    if "def " in code:
        parts = code.split("def ")
        if len(parts) > 1:
            code = "def " + parts[1]
    
    return '\n'.join(lines)

def compare_all_models(num_problems=164):
    """Evaluate and compare all post-trained models against baseline and self-eval"""
    model_results = {}
    percentages = [0.1, 1, 5, 10, 30, 50, 100]  
    
    # Get baseline results
    try:
        with open('results/baseline_results.json', 'r') as f:
            baseline_results = json.load(f)
            baseline_success = sum(1 for r in baseline_results if r.get('execution_success', False))
            baseline_total = len(baseline_results)
            baseline_rate = baseline_success / baseline_total * 100
            model_results['baseline'] = {
                'success': baseline_success,
                'total': baseline_total,
                'rate': baseline_rate
            }
            print(f"\nBaseline: {baseline_success}/{baseline_total} ({baseline_rate:.2f}%)")
    except (FileNotFoundError, json.JSONDecodeError):
        print("No baseline results found")
    
    # Get self-eval results
    try:
        with open('results/self_eval_results.json', 'r') as f:
            self_eval_results = json.load(f)
            self_eval_success = 0
            for r in self_eval_results:
                if r.get('improved_success', False) or r.get('execution_success', False):
                    self_eval_success += 1
            self_eval_total = len(self_eval_results)
            self_eval_rate = self_eval_success / self_eval_total * 100
            model_results['self-eval'] = {
                'success': self_eval_success,
                'total': self_eval_total,
                'rate': self_eval_rate
            }
            print(f"Self-evaluated: {self_eval_success}/{self_eval_total} ({self_eval_rate:.2f}%)")
    except (FileNotFoundError, json.JSONDecodeError):
        print("No self-evaluation results found")
    
    # Load existing results for each post-trained model instead of re-evaluating
    for percentage in percentages:
        results_path = f'results/post_trained_{percentage}pct/humaneval_results.json'
        try:
            with open(results_path, 'r') as f:
                results = json.load(f)
                success = sum(1 for r in results if r.get('execution_success', False))
                total = len(results)
                rate = success / total * 100
                model_results[f'post-trained-{percentage}%'] = {
                    'success': success,
                    'total': total,
                    'rate': rate
                }
                print(f"Post-trained ({percentage}%): {success}/{total} ({rate:.2f}%)")
        except (FileNotFoundError, json.JSONDecodeError):
            print(f"No results found for {percentage}% post-trained model")
    
    # Generate visualizations
    visualize_results(model_results)
    
    return model_results

def visualize_results(results):
    """Create visualizations for model comparisons"""
    # Prepare data for plotting
    labels = []
    success_rates = []
    
    # Add baseline and self-eval if available
    if 'baseline' in results:
        labels.append('Baseline')
        success_rates.append(results['baseline']['rate'])
    if 'self-eval' in results:
        labels.append('Self-Eval')
        success_rates.append(results['self-eval']['rate'])
    
    # Add post-trained models
    for key in sorted([k for k in results.keys() if k.startswith('post-trained')]):
        labels.append(key)
        success_rates.append(results[key]['rate'])
    
    # Add combined models if available
    for key in sorted([k for k in results.keys() if k.startswith('combined')]):
        labels.append(key)
        success_rates.append(results[key]['rate'])
    
    # Create bar chart
    plt.figure(figsize=(12, 7))
    bars = plt.bar(labels, success_rates, color=['blue', 'green', 'red', 'orange', 'purple', 'brown', 'pink', 'gray'])
    
    # Add labels and title
    plt.xlabel('Approach')
    plt.ylabel('Success Rate (%)')
    plt.title('HumanEval Success Rates by Approach')
    
    # Add value labels on bars
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.5,
                 f'{height:.1f}%', ha='center', va='bottom')
    
    # Add grid for readability
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    
    # Improve layout
    plt.tight_layout()
    
    # Save figure
    plt.savefig('results/model_comparison.png')
    plt.savefig('results/model_comparison.pdf')
    print("Visualization saved to results/model_comparison.png and .pdf")
    
    # Show plot
    plt.show()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate post-trained model")
    parser.add_argument("--percentage", type=int, default=1, 
                      help="Percentage of data used for post-training")
    parser.add_argument("--no-compare", action="store_true",
                      help="Skip comparison after evaluation")
    parser.add_argument("--compare-only", action="store_true",
                      help="Only run comparison without evaluation")
    args = parser.parse_args()
    
    if args.compare_only:
        # Only run comparison
        compare_all_models()
    else:
        # Run evaluation
        evaluate_post_trained_model(percentage=args.percentage)
        
        # Run comparison unless --no-compare flag is set
        if not args.no_compare:
            compare_all_models()