#!/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()