190 lines
7.0 KiB
Python
190 lines
7.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
多模型 benchmark 成本预测工具。
|
||
|
||
输入：
|
||
- 模型价格表 YAML（默认 bash/case/model_pricing.yaml）
|
||
- benchmark token 统计 CSV/Excel（默认 results/P800_benchmark_cost_GLM52.csv）
|
||
|
||
输出：
|
||
- 每个 benchmark 在每个模型下的预测成本矩阵（CSV/Excel）
|
||
- 按预算筛选出的可测 benchmark 列表
|
||
|
||
计价公式：
|
||
cost = n_samples * (input_tokens_mean * input_price + output_tokens_mean * output_price) * discount / 1_000_000
|
||
"""
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import pandas as pd
|
||
import yaml
|
||
|
||
|
||
ROOT = Path(__file__).resolve().parent.parent
|
||
DEFAULT_PRICING = ROOT / "bash" / "case" / "model_pricing.yaml"
|
||
DEFAULT_BENCHMARK = ROOT / "results" / "P800_benchmark_cost_GLM52.csv"
|
||
|
||
|
||
def load_pricing(path: Path) -> dict:
|
||
with open(path, encoding="utf-8") as f:
|
||
data = yaml.safe_load(f)
|
||
return data.get("models", {})
|
||
|
||
|
||
def load_benchmark(path: Path, sheet: str = None) -> pd.DataFrame:
|
||
if not path.exists():
|
||
raise FileNotFoundError(path)
|
||
suffix = path.suffix.lower()
|
||
if suffix in (".xlsx", ".xls"):
|
||
if sheet:
|
||
return pd.read_excel(path, sheet_name=sheet)
|
||
xl = pd.ExcelFile(path)
|
||
return pd.read_excel(path, sheet_name=xl.sheet_names[0])
|
||
return pd.read_csv(path)
|
||
|
||
|
||
def detect_columns(df: pd.DataFrame) -> dict:
|
||
def find(*candidates):
|
||
for c in candidates:
|
||
col = next((x for x in df.columns if x.lower() == c.lower()), None)
|
||
if col:
|
||
return col
|
||
return None
|
||
|
||
return {
|
||
"bench": find("benchmark", "benchmark名称"),
|
||
"n_samples": find("n_samples", "总样本数"),
|
||
"input": find("input_tokens_mean", "输入tokens_mean", "输入tokens"),
|
||
"output": find("output_tokens_mean", "输出tokens_mean", "输出tokens"),
|
||
"category": find("分类", "category"),
|
||
}
|
||
|
||
|
||
def compute_cost_matrix(df: pd.DataFrame, pricing: dict) -> pd.DataFrame:
|
||
cols = detect_columns(df)
|
||
for k, v in cols.items():
|
||
if v is None and k in ("bench", "n_samples", "input", "output"):
|
||
raise ValueError(f"缺少必要列: {k}")
|
||
|
||
df = df.copy()
|
||
n = pd.to_numeric(df[cols["n_samples"]], errors="coerce").fillna(0)
|
||
in_tok = pd.to_numeric(df[cols["input"]], errors="coerce").fillna(0)
|
||
out_tok = pd.to_numeric(df[cols["output"]], errors="coerce").fillna(0)
|
||
|
||
rows = []
|
||
for idx, row in df.iterrows():
|
||
bench = row[cols["bench"]]
|
||
category = row.get(cols["category"], "") if cols["category"] else ""
|
||
ns = n.iloc[idx]
|
||
it = in_tok.iloc[idx]
|
||
ot = out_tok.iloc[idx]
|
||
entry = {
|
||
"分类": category,
|
||
"Benchmark": bench,
|
||
"总样本数": ns,
|
||
"输入tokens_mean": it,
|
||
"输出tokens_mean": ot,
|
||
}
|
||
for model, p in pricing.items():
|
||
in_price = float(p["input_price"]) * float(p["discount"]) / 1_000_000.0
|
||
out_price = float(p["output_price"]) * float(p["discount"]) / 1_000_000.0
|
||
cost = ns * (it * in_price + ot * out_price)
|
||
entry[model] = round(cost, 4)
|
||
rows.append(entry)
|
||
|
||
return pd.DataFrame(rows)
|
||
|
||
|
||
def build_summary(matrix: pd.DataFrame, pricing: dict) -> pd.DataFrame:
|
||
"""为每个模型生成按分类/总计的汇总。"""
|
||
model_cols = list(pricing.keys())
|
||
summaries = []
|
||
|
||
# 分类汇总
|
||
if "分类" in matrix.columns and matrix["分类"].notna().any():
|
||
cat_sum = matrix.groupby("分类")[model_cols].sum().reset_index()
|
||
cat_sum["Benchmark"] = ""
|
||
summaries.append(cat_sum)
|
||
|
||
# 总计
|
||
total = {m: matrix[m].sum() for m in model_cols}
|
||
total["分类"] = ""
|
||
total["Benchmark"] = "TOTAL"
|
||
summaries.append(pd.DataFrame([total]))
|
||
|
||
return pd.concat(summaries, ignore_index=True)
|
||
|
||
|
||
def filter_by_budget(matrix: pd.DataFrame, budget: float, pricing: dict) -> pd.DataFrame:
|
||
"""返回每个模型在 budget 内可测的 benchmark 列表。"""
|
||
model_cols = list(pricing.keys())
|
||
keep = matrix[model_cols].le(budget).any(axis=1)
|
||
filtered = matrix[keep].copy()
|
||
for m in model_cols:
|
||
filtered[f"{m}_affordable"] = filtered[m] <= budget
|
||
return filtered
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="多模型 benchmark 成本预测")
|
||
parser.add_argument("--pricing", "-p", type=Path, default=DEFAULT_PRICING,
|
||
help="模型价格 YAML 路径")
|
||
parser.add_argument("--benchmark", "-b", type=Path, default=DEFAULT_BENCHMARK,
|
||
help="benchmark token 统计 CSV/Excel")
|
||
parser.add_argument("--sheet", "-s", default=None, help="Excel sheet 名")
|
||
parser.add_argument("--output", "-o", type=Path, default=ROOT / "results" / "P800_benchmark_cost_all_models",
|
||
help="输出文件路径（不含扩展名）")
|
||
parser.add_argument("--budget", type=float, default=None,
|
||
help="预算上限（元），用于筛选可测 benchmark")
|
||
parser.add_argument("--budget-output", type=Path, default=None,
|
||
help="预算筛选结果输出路径（默认 output_path_budget.csv/xlsx）")
|
||
args = parser.parse_args()
|
||
|
||
pricing = load_pricing(args.pricing)
|
||
if not pricing:
|
||
print(f"ERROR: 价格表为空: {args.pricing}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
df = load_benchmark(args.benchmark, args.sheet)
|
||
matrix = compute_cost_matrix(df, pricing)
|
||
summary = build_summary(matrix, pricing)
|
||
|
||
out_csv = Path(str(args.output) + ".csv")
|
||
out_xlsx = Path(str(args.output) + ".xlsx")
|
||
out_csv.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 保存完整成本矩阵
|
||
matrix.to_csv(out_csv, index=False, encoding="utf-8-sig")
|
||
with pd.ExcelWriter(out_xlsx, engine="openpyxl") as writer:
|
||
matrix.to_excel(writer, sheet_name="cost_matrix", index=False)
|
||
summary.to_excel(writer, sheet_name="summary", index=False)
|
||
|
||
print(f"成本矩阵已保存: {out_csv}, {out_xlsx}")
|
||
|
||
# 终端摘要
|
||
print("\n=== 各模型总成本（元）===")
|
||
total_row = summary[summary["Benchmark"] == "TOTAL"]
|
||
for m in pricing:
|
||
val = total_row[m].values[0] if not total_row.empty else 0
|
||
print(f" {m}: {val:.2f}")
|
||
|
||
# 预算筛选
|
||
if args.budget is not None:
|
||
budget_out = args.budget_output or Path(str(args.output) + f"_budget_{args.budget:.0f}")
|
||
filtered = filter_by_budget(matrix, args.budget, pricing)
|
||
filtered.to_csv(Path(str(budget_out) + ".csv"), index=False, encoding="utf-8-sig")
|
||
with pd.ExcelWriter(Path(str(budget_out) + ".xlsx"), engine="openpyxl") as writer:
|
||
filtered.to_excel(writer, sheet_name="affordable", index=False)
|
||
|
||
print(f"\n=== 预算 {args.budget:.2f} 元内可测的 benchmark ===")
|
||
for m in pricing:
|
||
names = filtered[filtered[f"{m}_affordable"]]["Benchmark"].tolist()
|
||
print(f" {m}: {len(names)} 个 -> {', '.join(names[:10])}" + (" ..." if len(names) > 10 else ""))
|
||
print(f"预算筛选结果已保存: {budget_out}.csv/.xlsx")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|