80 lines
3.1 KiB
Python
80 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
为 P800模型能力评测结果 Excel 的每个 sheet 追加 API 成本列。
|
|
只处理包含 Benchmark + 总样本数 + 输入/输出 tokens_mean 的 sheet。
|
|
"""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
|
|
def try_compute_cost(df: pd.DataFrame, input_price: float, output_price: float,
|
|
discount: float, model_name: str):
|
|
col_n = next((c for c in df.columns if '样本' in c and '数' in c), None)
|
|
col_in = next((c for c in df.columns if '输入tokens' in c), None)
|
|
col_out = next((c for c in df.columns if '输出tokens' in c), None)
|
|
col_bench = next((c for c in df.columns if 'benchmark' in c.lower()), None)
|
|
|
|
if col_n is None or col_in is None or col_out is None or col_bench is None:
|
|
return None
|
|
|
|
price_in = input_price * discount / 1_000_000.0
|
|
price_out = output_price * discount / 1_000_000.0
|
|
|
|
n = pd.to_numeric(df[col_n], errors='coerce').fillna(0)
|
|
in_tok = pd.to_numeric(df[col_in], errors='coerce').fillna(0)
|
|
out_tok = pd.to_numeric(df[col_out], errors='coerce').fillna(0)
|
|
|
|
result = df.copy()
|
|
result['input_cost_元'] = n * in_tok * price_in
|
|
result['output_cost_元'] = n * out_tok * price_out
|
|
result['total_cost_元'] = result['input_cost_元'] + result['output_cost_元']
|
|
result['price_model'] = model_name
|
|
return result
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--input', '-i', required=True)
|
|
parser.add_argument('--output', '-o', required=True)
|
|
parser.add_argument('--input-price', type=float, default=8.0)
|
|
parser.add_argument('--output-price', type=float, default=28.0)
|
|
parser.add_argument('--discount', '-d', type=float, default=0.65)
|
|
parser.add_argument('--model-name', '-m', default='GLM-5.2')
|
|
args = parser.parse_args()
|
|
|
|
xls = pd.ExcelFile(args.input)
|
|
out_path = Path(args.output)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with pd.ExcelWriter(out_path, engine='openpyxl') as writer:
|
|
for sheet in xls.sheet_names:
|
|
df = pd.read_excel(args.input, sheet_name=sheet)
|
|
new_df = try_compute_cost(
|
|
df, args.input_price, args.output_price,
|
|
args.discount, args.model_name,
|
|
)
|
|
if new_df is not None:
|
|
# 汇总
|
|
total_mask = new_df.iloc[:, 0].astype(str).str.lower().isin(['total', '总计', '合计'])
|
|
if not total_mask.any():
|
|
total_row = {'price_model': args.model_name}
|
|
for c in new_df.columns:
|
|
if 'cost_元' in c:
|
|
total_row[c] = new_df[c].sum()
|
|
total_df = pd.DataFrame([total_row])
|
|
new_df = pd.concat([new_df, total_df], ignore_index=True)
|
|
sheet_name = f"{sheet}_cost" if len(sheet) <= 25 else sheet[:25]
|
|
new_df.to_excel(writer, sheet_name=sheet_name, index=False)
|
|
else:
|
|
# 原样复制
|
|
df.to_excel(writer, sheet_name=sheet, index=False)
|
|
|
|
print(f"已保存: {out_path}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|