| |
| 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()) |
|
|
