reward-function/plot_capability.py at main · arpjw/reward-function · 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
"""Plot full-development capability measurements and their cost by mode."""
from __future__ import annotations

import argparse
import json
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--analysis', type=Path, required=True)
    parser.add_argument('--extension-analysis', type=Path)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    result = json.loads(args.analysis.read_text())
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    modes = result['diagnostic']['modes']
    labels = ['Direct\nanswer', '1024-token\nceiling']
    colors = ['#697586', '#1878b5']
    fig, axes = plt.subplots(1, 2, figsize=(9, 4), layout='constrained')
    accuracy = [100 * modes[m]['success'] for m in ('low', 'high')]
    cost = [modes[m]['cost'] for m in ('low', 'high')]
    if args.extension_analysis:
        extension = json.loads(args.extension_analysis.read_text())
        if any(result[k] != extension[k] for k in ('task_ids', 'task_sha256')):
            raise ValueError('extension uses different tasks')
        for key in ('model', 'revision', 'temperature', 'prompt_style'):
            if result['config'][key] != extension['config'][key]:
                raise ValueError('extension changes model or prompt')
        labels.append('2048-token\nceiling'); colors.append('#21866b')
        accuracy.append(100 * extension['diagnostic']['modes']['high']['success'])
        cost.append(extension['diagnostic']['modes']['high']['cost'])
    axes[0].bar(labels, accuracy, color=colors, width=.55)
    axes[0].set(ylim=(0, 110), ylabel='Correct repairs (%)', title='Accuracy on development tasks')
    axes[0].axhline(25, color='#a7afb8', linestyle=':', label='Uniform guessing')
    axes[0].legend(frameon=False, loc='upper left')
    axes[1].bar(labels, cost, color=colors, width=.55)
    axes[1].set(ylim=(0, max(cost) * 1.18), ylabel='Mean generated tokens', title='Measured completion cost')
    for ax, values, suffix in ((axes[0], accuracy, '%'), (axes[1], cost, '')):
        for i, value in enumerate(values):
            ax.annotate(f'{value:.1f}{suffix}', (i, value), xytext=(0, 5), textcoords='offset points', ha='center')
        ax.spines[['top', 'right']].set_visible(False)
    fig.suptitle(f'Qwen3-4B, 4-bit: {result["diagnostic"]["task_count"]} development tasks\nModel and prompt selection, not an adaptive-reward result', fontsize=12)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output.with_suffix('.png'), dpi=180)
    fig.savefig(args.output.with_suffix('.svg'))
    plt.close(fig)


if __name__ == '__main__':
    main()