run_self_eval.py

from StarCoder import build_model, generate_code, execute
import json
import os
from datasets import load_dataset 

def generate_self_evaluation(model, tokenizer, failed_result):
    """Generate analysis of failed code"""
    # More specific and detailed prompt for better analysis
    eval_prompt = f"""Analyze this Python code implementation that failed testing:

FUNCTION DESCRIPTION:
{failed_result.get('prompt', 'No description available')}

GENERATED CODE:
{failed_result['generated_code']}

TEST FAILURE:
{failed_result['output']}

Provide a detailed analysis of:
1. Logic errors found
2. Implementation issues (be specific about what's wrong)
3. Edge cases missed
4. Specific code fixes needed with corrected implementation

Focus on substantive issues rather than style. Be concise and specific.
"""
    return generate_code(model, tokenizer, eval_prompt)

def preprocess_improved_code(code, task_id, problem):
    """Preprocess the improved code for better test success"""
    # Add imports for typing if needed
    if any(typ in code for typ in ["List", "Optional", "Tuple", "Any"]):
        if "from typing import" not in code:
            code = "from typing import List, Optional, Tuple, Any\n\n" + code
    
    # Extract function name from problem if possible
    if problem and 'prompt' in problem:
        import re
        func_match = re.search(r'def\s+([a-zA-Z_][a-zA-Z0-9_]*)', problem['prompt'])
        if func_match:
            func_name = func_match.group(1)
            # Ensure the function name is correct in the generated code
            if func_name in code and not code.strip().startswith(f"def {func_name}"):
                code_parts = code.split(f"def {func_name}")
                if len(code_parts) > 1:
                    code = f"def {func_name}" + code_parts[1]
    
    # Clean up code structure
    code_lines = [line for line in code.split('\n') 
                 if not line.strip().startswith(('assert', '>>>'))]
    
    return '\n'.join(code_lines)

def improve_with_self_eval(baseline_results_path='results/baseline_results.json'):
    """Improve code generation using self-evaluation"""
    model, tokenizer, _ = build_model()
    
    # Load baseline results
    with open(baseline_results_path) as f:
        baseline_results = json.load(f)
    
    # Load HumanEval dataset to get test cases
    dataset = load_dataset("openai_humaneval")
    problems = {p['task_id']: p for p in dataset['test']}
    
    improved_results = []
    for result in baseline_results:
        task_id = result['task_id']
        problem = problems.get(task_id)
        
        if not result['execution_success'] and problem:
            print(f"\nProcessing failed task: {task_id}")
            
            # Generate evaluation
            evaluation = generate_self_evaluation(model, tokenizer, result)
            print("Generated self-evaluation")
            
            # Try again with evaluation feedback - improved prompt
            improved_prompt = f"""
                Complete the following Python function. Fix the issues identified below:

                ORIGINAL PROMPT:
                {problem['prompt']}

                PREVIOUS IMPLEMENTATION ISSUES:
                {evaluation}

                Write a complete, correct implementation that fixes these issues.
                Make sure your solution has proper error handling and handles all edge cases.
                """
            improved_code = generate_code(model, tokenizer, improved_prompt)
            print("Generated improved code")
            
            # Process and clean the improved code
            improved_code = preprocess_improved_code(improved_code, task_id, problem)
            
            # Test improved version with actual test cases
            test_code = problem['test']  # Get test cases from HumanEval dataset
            success, output = execute(improved_code, test_code)
            
            improved_results.append({
                **result,
                'self_evaluation': evaluation,
                'improved_code': improved_code,
                'improved_success': success,
                'improved_output': output,
                'test_cases': test_code
            })
            
            print(f"Execution success: {success}")
        else:
            improved_results.append(result)
    
    # Save results
    os.makedirs('results', exist_ok=True)
    save_path = 'results/self_eval_results.json'
    with open(save_path, 'w') as f:
        json.dump(improved_results, f, indent=2)
    
    # Print summary
    initial_success = sum(1 for r in baseline_results if r['execution_success'])
    
    # Fixed success calculation
    improved_success = 0
    for r in improved_results:
        if r.get('improved_success', False) == True:  # New successes
            improved_success += 1
        elif r.get('execution_success', False) == True:  # Original successes
            improved_success += 1
    
    total = len(baseline_results)
    
    print(f"\nResults Summary:")
    print(f"Initial success rate: {initial_success}/{total} ({initial_success/total*100:.2f}%)")
    print(f"Improved success rate: {improved_success}/{total} ({improved_success/total*100:.2f}%)")
    print(f"Results saved to {save_path}")
    
    return improved_results

if __name__ == "__main__":
    improve_with_self_eval()