evaluate_saved_model_StarCoder.py

import torch
from StarCoder import build_model, load_model, evaluate_code, config
import matplotlib.pyplot as plt
import json
import os


def get_output_filename():
    """Generate output filename based on test_size."""
    # Convert test_size to percentage
    percentage = int(config.test_size * 100)
    return f"StarCoderOutput_{percentage}%.md"
    
def get_curve_filename():
    """Generate filename for learning curves plot based on test_size."""
    percentage = int(config.test_size * 100)
    return f"Learning_Curves_StarCoderBase_{percentage}%.png"

def plot_learning_curves(history=None):
    """Plot training vs. validation learning curves with dynamic filename."""
    # If history is not provided, try to load it from file
    if history is None:
        try:
            # Attempt to load history from saved JSON file
            history_path = f"./model_checkpoint_starcoderbase/training_history.json"
            if os.path.exists(history_path):
                with open(history_path, "r") as f:
                    history = json.load(f)
                print(f"Loaded training history from {history_path}")
            else:
                print(f"Warning: No history file found at {history_path}")
                return
        except Exception as e:
            print(f"Error loading training history: {e}")
            return
        plt.figure(figsize=(10, 6))
    epochs = range(1, len(history['train_loss']) + 1)
    plt.plot(epochs, history['train_loss'], label='Training Loss')
    plt.plot(epochs, history['valid_loss'], label='Validation Loss')
    plt.title(f'StarCoderBase Learning Curves (Training Data: {int(config.test_size*100)}%)')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    plt.xticks(epochs)
    
    # Use dynamic filename based on test_size
    filename = get_curve_filename()
    plt.savefig(filename)
    plt.close()
    
    print(f"Learning curves saved to {filename}")

def evaluate_saved_model():
    # Initialize model and tokenizer
    model, tokenizer, device = build_model()
    
    # Load saved weights
    saved_model_path = "./model_checkpoint_starcoderbase"
    try:
        model, tokenizer = load_model(saved_model_path)
        print(f"Successfully loaded model and tokenizer from {saved_model_path}")
    except Exception as e:
        print(f"Error loading model: {e}")
        return

    # Define test cases
    test_cases = [
        {
            'func_documentation_string': 'Add two numbers a and b',
            'expected_output': '12',
            'test_input': 'print(add(5,7))'
        },
        {
            'func_documentation_string': 'Calculate factorial of a number n',
            'expected_output': '120',
            'test_input': 'print(factorial(5))'
        },
        {
            'func_documentation_string': 'Check if a number is prime',
            'expected_output': 'True',
            'test_input': 'print(is_prime(17))'
        },
        {
            'func_documentation_string': 'Reverse a string',
            'expected_output': 'olleh',
            'test_input': 'print(reverse_string("hello"))'
        },
        {
            'func_documentation_string': 'Find the maximum value in a list of numbers',
            'expected_output': '42',
            'test_input': 'print(find_max([5, 42, 17, 8, 1]))'
        },
        {
            'func_documentation_string': 'Count the number of vowels in a string',
            'expected_output': '2',
            'test_input': 'print(count_vowels("hello"))'
        },
        {
            'func_documentation_string': 'Generate Fibonacci sequence up to n terms',
            'expected_output': '[0, 1, 1, 2, 3]',
            'test_input': 'print(fibonacci(5))'
        },
        {
            'func_documentation_string': 'Convert temperature from Celsius to Fahrenheit',
            'expected_output': '98.6',
            'test_input': 'print(celsius_to_fahrenheit(37))'
        }
    ]
    
    # Run evaluation
    print("\nEvaluating saved model...")
    evaluation_results = evaluate_code(model, tokenizer, test_cases)
    
    # Print results
    print("\nEvaluation Results:")
    print(f"Syntax Accuracy: {evaluation_results['syntax_accuracy']:.2f}%")
    print(f"Execution Accuracy: {evaluation_results['execution_accuracy']:.2f}%")
    print(f"Functional Accuracy: {evaluation_results['functional_accuracy']:.2f}%")

    # Write results to StarCoderOutput.md
    output_filename = get_output_filename()
    with open(output_filename, "w") as f:
        f.write(f"# StarCoderBase Model Evaluation Results (Training Data: {int(config.test_size*100)}%\n\n")
        
        # Summary statistics
        f.write("## Summary\n")
        f.write(f"- **Training Data Size:** {config.test_size*100:.0f}% of dataset\n")
        f.write(f"- **Syntax Accuracy:** {evaluation_results['syntax_accuracy']:.2f}%\n")
        f.write(f"- **Execution Accuracy:** {evaluation_results['execution_accuracy']:.2f}%\n")
        f.write(f"- **Functional Accuracy:** {evaluation_results['functional_accuracy']:.2f}%\n\n")
        
        # Individual test cases
        f.write("## Test Case Results\n\n")
        for i, (test_case, result) in enumerate(zip(test_cases, evaluation_results['test_results'])):
            f.write(f"### Test {i+1}: {test_case['func_documentation_string']}\n")
            f.write("```python\n")
            f.write(result['generated_code'])
            f.write("\n```\n\n")
            f.write(f"**Test Input:** `{test_case['test_input']}`\n")
            f.write(f"**Expected Output:** `{test_case['expected_output']}`\n")
            f.write(f"**Actual Output:** `{result['output']}`\n")
            f.write(f"**Success:** {result['success']}\n")
            f.write(f"**Output Matches:** {result['output_matches']}\n\n")
    
    print(f"Results saved to {output_filename}")

    plot_learning_curves()
if __name__ == "__main__":
    evaluate_saved_model()