NumberNames/plot_error_rates.py at main · SakanaAI/NumberNames · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
#!/usr/bin/env python3
"""
Plot error rates from dev_out.txt files.
"""
import os
import sys
import glob
import matplotlib.pyplot as plt
import numpy as np


def parse_dev_out_file(filepath):
    """
    Parse a dev_out.txt file and extract error rates from the last 3 lines.
    Returns a dict with 'string_error', 'word_error', and 'char_error'.
    """
    with open(filepath, 'r') as f:
        lines = f.readlines()

    # Get the last 3 lines (Names, Words, Chars)
    last_lines = [line.strip() for line in lines[-3:]]

    results = {}
    for line in last_lines:
        parts = line.split('\t')
        if len(parts) >= 4:
            metric = parts[0].rstrip(':')
            error_rate = float(parts[3])

            if metric == 'Names':
                results['string_error'] = error_rate
            elif metric == 'Words':
                results['word_error'] = error_rate
            elif metric == 'Chars':
                results['char_error'] = error_rate

    return results


def plot_error_rates(data_dir, output_file, dpi=300):
    """
    Plot error rates from all dev_out.txt files in the given directory.

    Args:
        data_dir: Directory containing the .dev_out.txt files
        output_file: Path to save the plot (should end in .pdf or .png)
        dpi: DPI for PNG output (default: 300), ignored for PDF
    """
    # Find all dev_out.txt files
    pattern = os.path.join(data_dir, '*.dev_out.txt')
    files = sorted(glob.glob(pattern))

    if not files:
        print(f"No .dev_out.txt files found in {data_dir}")
        return

    # Parse all files
    languages = []
    string_errors = []
    word_errors = []
    char_errors = []

    for filepath in files:
        # Extract language name from filename
        filename = os.path.basename(filepath)
        language = filename.replace('.dev_out.txt', '')

        # Parse error rates
        results = parse_dev_out_file(filepath)

        languages.append(language)
        string_errors.append(results['string_error'])
        word_errors.append(results['word_error'])
        char_errors.append(results['char_error'])

    # Create the plot
    x = np.arange(len(languages))
    width = 0.25

    fig, ax = plt.subplots(figsize=(12, 6))

    bars1 = ax.bar(x - width, string_errors, width, label='String Error Rate', alpha=0.8)
    bars2 = ax.bar(x, word_errors, width, label='Word Error Rate', alpha=0.8)
    bars3 = ax.bar(x + width, char_errors, width, label='Character Error Rate', alpha=0.8)

    ax.set_xlabel('Language', fontsize=12)
    ax.set_ylabel('Error Rate', fontsize=12)
    ax.set_title('Error Rates by Language', fontsize=14, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(languages, rotation=45, ha='right')
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    # Set y-axis to start at 0 and go to 1 (or slightly higher if needed)
    ax.set_ylim(0, max(max(string_errors), 1.0) * 1.1)

    plt.tight_layout()

    # Auto-detect format from file extension
    file_ext = os.path.splitext(output_file)[1].lower()
    if file_ext == '.png':
        plt.savefig(output_file, format='png', dpi=dpi, bbox_inches='tight')
    elif file_ext == '.pdf':
        plt.savefig(output_file, format='pdf', bbox_inches='tight')
    else:
        # Default to PDF if extension is unclear
        plt.savefig(output_file, format='pdf', bbox_inches='tight')

    print(f"Plot saved to {output_file}")

    # Print summary
    print(f"\nSummary for {data_dir}:")
    print(f"{'Language':<25} {'String ER':>10} {'Word ER':>10} {'Char ER':>10}")
    print("-" * 60)
    for i, lang in enumerate(languages):
        print(f"{lang:<25} {string_errors[i]:>10.2f} {word_errors[i]:>10.2f} {char_errors[i]:>10.2f}")


if __name__ == '__main__':
    if len(sys.argv) > 1:
        # If directory is provided as argument
        data_dir = sys.argv[1]
        # Check if output format is specified (default to PDF)
        output_format = sys.argv[2] if len(sys.argv) > 2 else 'pdf'
        output_file = os.path.join(data_dir, f'plot.{output_format}')
        plot_error_rates(data_dir, output_file)
    else:
        # Default: process both directories, generate both PDF and PNG
        print("Processing coding task results...")
        coding_dir = 'tasks/coding/dev_out_data'
        plot_error_rates(coding_dir, os.path.join(coding_dir, 'plot.pdf'))
        plot_error_rates(coding_dir, os.path.join(coding_dir, 'plot.png'))

        print("\n" + "="*60 + "\n")

        print("Processing reasoning task results...")
        reasoning_dir = 'tasks/reasoning/dev_out_data'
        plot_error_rates(reasoning_dir, os.path.join(reasoning_dir, 'plot.pdf'))
        plot_error_rates(reasoning_dir, os.path.join(reasoning_dir, 'plot.png'))