138 lines
5.4 KiB
Python
138 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
根据 benchmark 结果计算 API 调用成本。
|
||
|
||
支持输入：
|
||
- CSV（如 P800模型能力评测结果 - DS4-Flash-INT8-NO-Thinking-2.0-FULL.csv）
|
||
- Excel（如 P800模型能力评测结果_统一格式_filled.xlsx）
|
||
|
||
计价公式：
|
||
input_cost = n_samples * input_tokens_mean / 1_000_000 * input_price * discount
|
||
output_cost = n_samples * output_tokens_mean / 1_000_000 * output_price * discount
|
||
|
||
价格单位：元 / 百万 tokens（按国内 API 常见报价）。
|
||
"""
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import pandas as pd
|
||
|
||
|
||
def load_input(path: str, sheet: str = None) -> pd.DataFrame:
|
||
p = Path(path)
|
||
if not p.exists():
|
||
raise FileNotFoundError(path)
|
||
|
||
if p.suffix.lower() in ('.xlsx', '.xls'):
|
||
if sheet:
|
||
return pd.read_excel(path, sheet_name=sheet)
|
||
xl = pd.ExcelFile(path)
|
||
print(f"可用 sheet: {xl.sheet_names}")
|
||
return pd.read_excel(path, sheet_name=xl.sheet_names[0])
|
||
else:
|
||
return pd.read_csv(path)
|
||
|
||
|
||
def compute_cost(
|
||
df: pd.DataFrame,
|
||
input_price: float,
|
||
output_price: float,
|
||
discount: float,
|
||
model_name: str,
|
||
) -> pd.DataFrame:
|
||
df = df.copy()
|
||
|
||
# 兼容中英文列名
|
||
col_n_samples = next((c for c in df.columns if '样本' in c or c.lower() == 'n_samples'), None)
|
||
col_input = next((c for c in df.columns if '输入tokens' in c or c.lower() == 'input_tokens_mean'), None)
|
||
col_output = next((c for c in df.columns if '输出tokens' in c or c.lower() == 'output_tokens_mean'), None)
|
||
col_total = next((c for c in df.columns if '累计总tokens' in c or c.lower() == 'total_tokens'), None)
|
||
col_score = next((c for c in df.columns if '得分' in c or c.lower() == 'score'), None)
|
||
col_bench = next((c for c in df.columns if 'benchmark' in c.lower()), None)
|
||
|
||
if col_n_samples is None or col_input is None or col_output is None:
|
||
raise ValueError("输入文件缺少必要列：总样本数 / 输入tokens_mean / 输出tokens_mean")
|
||
|
||
price_in = input_price * discount / 1_000_000.0
|
||
price_out = output_price * discount / 1_000_000.0
|
||
|
||
# 数值化，空值填 0
|
||
n = pd.to_numeric(df[col_n_samples], errors='coerce').fillna(0)
|
||
in_tok = pd.to_numeric(df[col_input], errors='coerce').fillna(0)
|
||
out_tok = pd.to_numeric(df[col_output], errors='coerce').fillna(0)
|
||
|
||
df['input_tokens_total'] = (n * in_tok).astype('int64')
|
||
df['output_tokens_total'] = (n * out_tok).astype('int64')
|
||
df['input_cost_yuan'] = n * in_tok * price_in
|
||
df['output_cost_yuan'] = n * out_tok * price_out
|
||
df['total_cost_yuan'] = df['input_cost_yuan'] + df['output_cost_yuan']
|
||
|
||
# 元/千 tokens 便于查看
|
||
df['input_price_per_1k_tokens'] = price_in * 1000
|
||
df['output_price_per_1k_tokens'] = price_out * 1000
|
||
df['model_pricing'] = f"{model_name} (in={input_price}*{-discount}, out={output_price}*{discount})"
|
||
|
||
return df
|
||
|
||
|
||
def save_output(df: pd.DataFrame, path: str):
|
||
p = Path(path)
|
||
p.parent.mkdir(parents=True, exist_ok=True)
|
||
if p.suffix.lower() in ('.xlsx', '.xls'):
|
||
with pd.ExcelWriter(path, engine='openpyxl') as writer:
|
||
df.to_excel(writer, sheet_name='cost', index=False)
|
||
else:
|
||
df.to_csv(path, index=False, encoding='utf-8-sig')
|
||
print(f"结果已保存: {path}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="根据 benchmark 结果计算 API 成本")
|
||
parser.add_argument('--input', '-i', required=True, help='输入 CSV/Excel 路径')
|
||
parser.add_argument('--sheet', '-s', default=None, help='Excel sheet 名（默认第一个）')
|
||
parser.add_argument('--input-price', type=float, default=8.0, help='输入单价（元/百万 tokens）')
|
||
parser.add_argument('--output-price', type=float, default=28.0, help='输出单价（元/百万 tokens）')
|
||
parser.add_argument('--discount', '-d', type=float, default=0.65, help='折扣率，如 0.65')
|
||
parser.add_argument('--model-name', '-m', default='GLM-5.2', help='模型名称（仅用于标注）')
|
||
parser.add_argument('--output', '-o', required=True, help='输出 CSV/Excel 路径')
|
||
args = parser.parse_args()
|
||
|
||
df = load_input(args.input, args.sheet)
|
||
df_cost = compute_cost(
|
||
df,
|
||
input_price=args.input_price,
|
||
output_price=args.output_price,
|
||
discount=args.discount,
|
||
model_name=args.model_name,
|
||
)
|
||
|
||
# 汇总行（如果原表没有 totals，追加一个）
|
||
numeric_cols = ['input_cost_yuan', 'output_cost_yuan', 'total_cost_yuan']
|
||
total_row = {c: df_cost[c].sum() for c in numeric_cols}
|
||
total_row['Benchmark'] = 'TOTAL'
|
||
total_row['model_pricing'] = args.model_name
|
||
|
||
# 插入或更新 TOTAL 行
|
||
is_total = df_cost['Benchmark'].astype(str).str.lower().isin(['total', '总计', '合计'])
|
||
if is_total.any():
|
||
for c in numeric_cols:
|
||
df_cost.loc[is_total, c] = total_row[c]
|
||
else:
|
||
total_df = pd.DataFrame([total_row])
|
||
df_cost = pd.concat([df_cost, total_df], ignore_index=True)
|
||
|
||
save_output(df_cost, args.output)
|
||
|
||
# 终端摘要
|
||
summary = df_cost[['Benchmark', '总样本数', '输入tokens_mean', '输出tokens_mean',
|
||
'input_cost_yuan', 'output_cost_yuan', 'total_cost_yuan']].copy()
|
||
print('\n=== 成本摘要（元）===')
|
||
print(summary.to_string(index=False))
|
||
print(f"\n总成本: {total_row['total_cost_yuan']:.4f} 元")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|