File size: 2,995 Bytes
c6856c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
from pathlib import Path
def load_jsonl(path: Path) -> list[dict]:
rows = []
with path.open("r", encoding="utf-8", errors="replace") as handle:
for raw_line in handle:
line = raw_line.strip()
if not line:
continue
try:
payload = json.loads(line)
except json.JSONDecodeError:
continue
if isinstance(payload, dict):
rows.append(payload)
return rows
def as_int(value) -> int:
if isinstance(value, bool):
return int(value)
if isinstance(value, int):
return value
if isinstance(value, float):
return int(value)
if isinstance(value, str) and value.strip():
try:
return int(float(value))
except ValueError:
return 0
return 0
def summarize(task_name: str, path: Path) -> dict:
rows = load_jsonl(path)
prompt_vals = [as_int(row.get("prompt_tokens")) for row in rows]
completion_vals = [as_int(row.get("completion_tokens")) for row in rows]
total_vals = [as_int(row.get("total_tokens")) for row in rows]
call_vals = [as_int(row.get("llm_call_count")) for row in rows]
with_usage = sum(1 for row in rows if any(as_int(row.get(key)) > 0 for key in ("prompt_tokens", "completion_tokens", "total_tokens")))
n = len(rows)
return {
"task": task_name,
"cases": n,
"cases_with_usage": with_usage,
"total_prompt_tokens": sum(prompt_vals),
"total_completion_tokens": sum(completion_vals),
"total_tokens": sum(total_vals),
"avg_prompt_tokens_per_case": round(sum(prompt_vals) / n, 2) if n else 0.0,
"avg_completion_tokens_per_case": round(sum(completion_vals) / n, 2) if n else 0.0,
"avg_total_tokens_per_case": round(sum(total_vals) / n, 2) if n else 0.0,
"avg_llm_calls_per_case": round(sum(call_vals) / n, 2) if n else 0.0,
}
def main() -> int:
parser = argparse.ArgumentParser(description="Summarize real token usage from LAB-Bench compact result JSONL files.")
parser.add_argument("--results-dir", type=Path, required=True)
args = parser.parse_args()
mapping = {
"DbQA": [
args.results_dir / "dbqa_results.jsonl",
args.results_dir / "dbqa_batch" / "dbqa_results.jsonl",
],
"SeqQA": [
args.results_dir / "seqqa_results.jsonl",
args.results_dir / "seqqa_batch" / "seqqa_results.jsonl",
],
}
summaries = []
for task, candidates in mapping.items():
path = next((candidate for candidate in candidates if candidate.exists()), None)
if path is None:
continue
summaries.append(summarize(task, path))
print(json.dumps(summaries, indent=2, ensure_ascii=False))
return 0
if __name__ == "__main__":
raise SystemExit(main())
|