diff --git a/benchmark_chart.py b/benchmark_chart.py new file mode 100644 index 0000000..cf2843e --- /dev/null +++ b/benchmark_chart.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +""" +AI 模型基准测试对比图生成器 + +用法: + python benchmark_chart.py -c data.py -o output.png + python benchmark_chart.py -c data.py -o output.png --dpi 150 + python benchmark_chart.py -c data.yaml -o result.png --width 28 --height 18 + +数据文件格式(支持 .py / .json / .yaml): + + TITLE = "总标题" + COLORS = {"qwen36": "#5B2D8E", ...} + BENCHMARKS = [ + { + "title": "Benchmark Name", + "subtitle": "简短描述", + "models": [ + {"name": "全名", "short_label": "短标", "group": "moe|dense", "score": 85.0, "color_key": "qwen36"}, + ... + ], + }, + ... + ] +""" + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import matplotlib +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker +import numpy as np + +matplotlib.use("Agg") + +# ────────────────────── 默认值 ────────────────────── +DEFAULT_COLORS = { + "qwen36": "#5B2D8E", + "qwen35": "#8B5CF6", + "gemma26": "#06B6D4", + "qwen27": "#B0B5BD", + "gemma31": "#34D399", +} +DEFAULT_BG = "#E8F5E9" +DEFAULT_BAR_WIDTH = 0.55 +DEFAULT_DPI = 200 +DEFAULT_FIGSIZE = (24, 16) +GROUP_GAP = 1.2 # MoE 和 Dense 组间距 + + +def load_data(source: str) -> dict: + p = Path(source) + if p.suffix == ".py": + sys.path.insert(0, str(p.parent)) + ns: dict[str, Any] = {} + exec(p.read_text(encoding="utf-8"), ns) + return ns + elif p.suffix == ".json": + return json.loads(p.read_text(encoding="utf-8")) + elif p.suffix in (".yaml", ".yml"): + try: + import yaml + except ImportError: + raise ImportError("需要 PyYAML: pip install pyyaml") + return yaml.safe_load(p.read_text(encoding="utf-8")) + else: + raise ValueError(f"不支持的文件格式: {p.suffix}") + + +def draw_chart(ax: plt.Axes, benchmark: dict, colors: dict, bar_width: float): + """在单个 Axes 上绘制一个基准测试的柱状图""" + title = benchmark.get("title", "") + subtitle = benchmark.get("subtitle", "") + models = benchmark.get("models", []) + + moe_models = [m for m in models if m.get("group", "").lower() == "moe"] + dense_models = [m for m in models if m.get("group", "").lower() == "dense"] + n_moe = len(moe_models) + n_dense = len(dense_models) + + # ── X 坐标 ── + moe_x = np.arange(n_moe) - (n_moe - 1) / 2 if n_moe else np.array([]) + dense_start = (n_moe - 1) / 2 + GROUP_GAP + bar_width / 2 if n_moe else 0 + dense_x = (np.arange(n_dense) + dense_start) if n_dense else np.array([]) + sep_x = (moe_x[-1] + dense_x[0]) / 2 if n_moe > 0 and n_dense > 0 else None + + # ── 画柱子 ── + all_scores: list[float] = [] + for i, m in enumerate(moe_models): + score = float(m.get("score", 0)) + all_scores.append(score) + ck = m.get("color_key", "") + color = colors.get(ck, "#999") + ax.bar(moe_x[i], score, bar_width, color=color, + edgecolor="white", linewidth=0.5, zorder=3) + for i, m in enumerate(dense_models): + score = float(m.get("score", 0)) + all_scores.append(score) + ck = m.get("color_key", "") + color = colors.get(ck, "#999") + ax.bar(dense_x[i], score, bar_width, color=color, + edgecolor="white", linewidth=0.5, zorder=3) + + if not all_scores: + return + + y_max = max(all_scores) + y_pad = y_max * 0.32 # 上方留白给分数标注 + 分组标签 + label_reserve = y_max * 0.22 # 底部留给 x 轴标签的空间 + + # ═══ X 轴标签:模型短名 ═══ + label_fs = 7.5 if n_moe + n_dense > 4 else 8.5 + name_y = -label_reserve * 0.4 + + for i, m in enumerate(moe_models): + lbl = m.get("short_label", m.get("name", "")) + ax.text(moe_x[i], name_y, lbl, ha="center", va="top", + fontsize=label_fs, color="#555", linespacing=1.0) + for i, m in enumerate(dense_models): + lbl = m.get("short_label", m.get("name", "")) + ax.text(dense_x[i], name_y, lbl, ha="center", va="top", + fontsize=label_fs, color="#555", linespacing=1.0) + + # ═══ MoE / Dense 分组标签:放在柱子组上方 ═══ + if n_moe > 0: + ax.text(moe_x.mean(), y_max + y_pad * 0.55, "MoE", + ha="center", va="center", + fontsize=8.5, fontweight="bold", color="#888", + bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#ddd", alpha=0.85)) + if n_dense > 0: + ax.text(dense_x.mean(), y_max + y_pad * 0.55, "Dense", + ha="center", va="center", + fontsize=8.5, fontweight="bold", color="#888", + bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#ddd", alpha=0.85)) + + # ── 虚线分隔 ── + if sep_x is not None: + ax.axvline(x=sep_x, ymin=0.04, ymax=0.92, color="#ccc", + linestyle=":", linewidth=1.2, zorder=2) + + # ── 最高分标注 ── + top_score = max(all_scores) + for i, m in enumerate(moe_models): + if m.get("score", 0) == top_score: + ax.annotate( + f"{top_score}", + xy=(moe_x[i], top_score), + xytext=(0, 5), textcoords="offset points", + ha="center", va="bottom", + fontsize=11, fontweight="bold", color="#222", + ) + break + else: + for i, m in enumerate(dense_models): + if m.get("score", 0) == top_score: + ax.annotate( + f"{top_score}", + xy=(dense_x[i], top_score), + xytext=(0, 5), textcoords="offset points", + ha="center", va="bottom", + fontsize=11, fontweight="bold", color="#222", + ) + break + + # ═══ 标题 + 副标题(放在上方,无重叠) ═══ + # 标题:左上角,加粗 + ax.text(0.0, 1.02, title, transform=ax.transAxes, + fontsize=12, fontweight="bold", color="#222", + ha="left", va="bottom") + if subtitle: + # 副标题:标题下方 + ax.text(0.0, 0.94, subtitle, transform=ax.transAxes, + fontsize=7.5, color="#aaa", ha="left", va="top", + fontstyle="italic") + + # ── 美化 ── + ax.set_facecolor("white") + for spine in ("top", "right", "left"): + ax.spines[spine].set_visible(False) + ax.spines["bottom"].set_color("#ddd") + ax.tick_params(axis="x", bottom=False, labelbottom=False) + ax.tick_params(axis="y", labelsize=6.5, colors="#aaa", length=0) + ax.yaxis.set_major_locator(mticker.MaxNLocator(4, integer=True)) + ax.grid(axis="y", color="#f0f0f0", linewidth=0.5, zorder=0) + + # ── 坐标范围 ── + all_x = list(moe_x) + list(dense_x) + x_pad = bar_width * 1.8 + ax.set_xlim(min(all_x) - x_pad, max(all_x) + x_pad) + ax.set_ylim(-label_reserve, y_max + y_pad) + + +def make_chart( + data_source: str, + output_path: str, + dpi: int = DEFAULT_DPI, + figsize: tuple = DEFAULT_FIGSIZE, +): + data = load_data(data_source) + benchmarks = data.get("benchmarks", data.get("BENCHMARKS", [])) + colors = data.get("colors", data.get("COLORS", DEFAULT_COLORS)) + bg_color = data.get("bg_color", data.get("BG_COLOR", DEFAULT_BG)) + bar_width = data.get("bar_width", data.get("BAR_WIDTH", DEFAULT_BAR_WIDTH)) + title = data.get("title", data.get("TITLE", "")) + ncols = data.get("ncols", data.get("NCOLS", 4)) + + n = len(benchmarks) + if n == 0: + print("❌ 没有找到 benchmarks 数据") + sys.exit(1) + + nrows = (n + ncols - 1) // ncols + + fig, axes = plt.subplots( + nrows, ncols, figsize=figsize, + facecolor=bg_color, + ) + fig.subplots_adjust( + hspace=0.55, wspace=0.35, + top=0.93, bottom=0.09, left=0.04, right=0.96, + ) + + # 统一 axes 数组 + axes_flat = np.array(axes).flatten() + + if title: + fig.suptitle(title, fontsize=16, fontweight="bold", color="#222", y=0.97) + + for i, bench in enumerate(benchmarks): + draw_chart(axes_flat[i], bench, colors, bar_width) + + for j in range(i + 1, len(axes_flat)): + axes_flat[j].set_visible(False) + + # ── 底部图例 ── + seen: dict[str, dict] = {} + for bench in benchmarks: + for m in bench.get("models", []): + name = m.get("name", "") + if name not in seen: + seen[name] = m + + if seen: + handles = [] + labels = [] + for name, m in seen.items(): + ck = m.get("color_key", "") + color = colors.get(ck, "#999") + handles.append(plt.Rectangle((0, 0), 1, 1, fc=color, ec="white", lw=0.5)) + labels.append(name) + fig.legend( + handles, labels, loc="lower center", + ncol=min(len(seen), 5), frameon=False, + fontsize=8, + bbox_to_anchor=(0.5, 0.005), + ) + + fig.savefig(output_path, dpi=dpi, bbox_inches="tight", facecolor=bg_color) + print(f"✅ 图表已保存到 {output_path}") + return output_path + + +def main(): + parser = argparse.ArgumentParser(description="AI Benchmark 对比图生成器") + parser.add_argument("-c", "--config", required=True, + help="数据配置文件 (.py / .json / .yaml)") + parser.add_argument("-o", "--output", default="benchmark_chart.png", + help="输出图片路径") + parser.add_argument("--dpi", type=int, default=DEFAULT_DPI, + help=f"分辨率 (默认: {DEFAULT_DPI})") + parser.add_argument("--width", type=float, default=DEFAULT_FIGSIZE[0], + help="图宽 (英寸)") + parser.add_argument("--height", type=float, default=DEFAULT_FIGSIZE[1], + help="图高 (英寸)") + args = parser.parse_args() + make_chart( + data_source=args.config, + output_path=args.output, + dpi=args.dpi, + figsize=(args.width, args.height), + ) + + +if __name__ == "__main__": + main()