huangzhuang_3rd/benchmark-chart
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add benchmark_chart.py

Browse Source
This commit is contained in:
huangzhuang_3rd
2026-06-02 00:44:51 +08:00
parent dd89d915a9
commit 9746adc05a
1 changed files with 289 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

289
benchmark_chart.py Normal file
View File

@@ -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()