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