File size: 4,536 Bytes
e2bfccc | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from __future__ import annotations
import argparse
import csv
import json
from pathlib import Path
from typing import Any
def _as_float(value: str | None) -> float | None:
if value is None or value == "":
return None
try:
return float(value)
except ValueError:
return None
def _load_rows(root: Path) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
for csv_path in sorted(root.glob("*/taonet_real_token_benchmark.csv")):
variant = csv_path.parent.name
with csv_path.open("r", newline="", encoding="utf-8") as handle:
for row in csv.DictReader(handle):
row = dict(row)
row["variant"] = variant
rows.append(row)
return rows
def _best_forward_backward(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
candidates = [row for row in rows if row.get("mode") == "forward_backward"]
grouped: dict[str, list[dict[str, Any]]] = {}
for row in candidates:
grouped.setdefault(row["variant"], []).append(row)
best_rows = []
for variant, items in grouped.items():
items.sort(
key=lambda row: (
_as_float(row.get("eval_loss")) if _as_float(row.get("eval_loss")) is not None else float("inf"),
-(_as_float(row.get("eval_accuracy")) or 0.0),
)
)
best_rows.append(items[0])
best_rows.sort(
key=lambda row: (
_as_float(row.get("eval_loss")) if _as_float(row.get("eval_loss")) is not None else float("inf"),
-(_as_float(row.get("eval_accuracy")) or 0.0),
)
)
return best_rows
def _project(row: dict[str, Any]) -> dict[str, Any]:
keys = [
"variant",
"architecture",
"hybrid_pattern",
"batch_size",
"seq_len",
"total_params",
"ssm_core",
"ssm_hidden_dim",
"ssm_mixer_dim",
"ssm_num_lanes",
"ssm_lane_mode",
"ssm_split_mix",
"tokens_per_s_mean",
"eval_loss",
"eval_perplexity",
"eval_accuracy",
"train_final_loss",
"train_seconds",
"peak_reserved_mb",
"case_id",
"checkpoint_path",
]
return {key: row.get(key, "") for key in keys}
def _write_markdown(summary: list[dict[str, Any]], path: Path) -> None:
headers = [
"variant",
"architecture",
"batch",
"params",
"eval_loss",
"eval_acc",
"tok/s",
"checkpoint",
]
lines = [
"# TaoNet Benchmark Suite Summary",
"",
"| " + " | ".join(headers) + " |",
"| " + " | ".join(["---"] * len(headers)) + " |",
]
for row in summary:
lines.append(
"| "
+ " | ".join(
[
str(row["variant"]),
str(row["architecture"]),
str(row["batch_size"]),
str(row["total_params"]),
str(row["eval_loss"]),
str(row["eval_accuracy"]),
str(row["tokens_per_s_mean"]),
str(row["checkpoint_path"]),
]
)
+ " |"
)
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
def main() -> None:
parser = argparse.ArgumentParser(description="Summarize a TaoNet benchmark suite output directory.")
parser.add_argument("--suite-dir", required=True, help="Directory containing one subdirectory per benchmark variant.")
parser.add_argument("--output-json", default="", help="Summary JSON path. Defaults to <suite-dir>/suite_summary.json.")
parser.add_argument("--output-md", default="", help="Summary Markdown path. Defaults to <suite-dir>/suite_summary.md.")
args = parser.parse_args()
suite_dir = Path(args.suite_dir)
rows = _load_rows(suite_dir)
summary = [_project(row) for row in _best_forward_backward(rows)]
json_path = Path(args.output_json) if args.output_json else suite_dir / "suite_summary.json"
md_path = Path(args.output_md) if args.output_md else suite_dir / "suite_summary.md"
json_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
_write_markdown(summary, md_path)
print(f"Wrote {json_path}")
print(f"Wrote {md_path}")
if __name__ == "__main__":
main()
|