290 lines
10 KiB
Python
290 lines
10 KiB
Python
#!/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()
|