Script to create an aggregate summary table comparing useful metrics for Pruna vs Base (HQQ+torch.compile)

aggregate_values_benchmark_table.py

"""
Aggregate Summary Table — Pruna vs Base(HQQ+torch.compile)
=======================================

Loads raw benchmark CSVs, computes aggregate_summary.csv,
then saves it as a styled PNG table.

Input
-----
benchmark/raw_benchmark_results_hqq.csv
benchmark/raw_benchmark_results_pruna.csv

Outputs
-------
benchmark/eda_outputs/
├── aggregate_summary.csv
└── aggregate_summary_table_styled.png
"""

from pathlib import Path

import pandas as pd
import matplotlib.pyplot as plt


# ============================================================
# CONFIG
# ============================================================

HQQ_PATH   = "/workspaces/pruna-cookbook/benchmark/raw_benchmark_results_hqq.csv"
PRUNA_PATH = "/workspaces/pruna-cookbook/benchmark/raw_benchmark_results_pruna.csv"

OUTPUT_DIR = Path("/workspaces/pruna-cookbook/benchmark/eda_outputs")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

FAILED_THRESHOLD = 1000

# ============================================================
# DESIGN SYSTEM
# ============================================================

PRUNA_COLOR  = "#6A0DAD"   # purple
HQQ_COLOR    = "#E05C2A"   # amber-red
PRUNA_LIGHT  = "#F3E8FF"
HQQ_LIGHT    = "#FFF0E8"
HEADER_BG    = "#1A1A2E"
ROW_LABEL_BG = "#F5F5F5"

# ============================================================
# LOAD, CLEAN & MERGE
# ============================================================

print("Loading data...")

hqq_df        = pd.read_csv(HQQ_PATH)
pruna_df       = pd.read_csv(PRUNA_PATH)
hqq_df["framework"]   = "HQQ"
pruna_df["framework"]  = "Pruna"

df = pd.concat([hqq_df, pruna_df], ignore_index=True)

numeric_cols = [
    "prefill_latency_s",
    "decode_tokens_per_sec",
    "avg_decode_latency_per_token_ms",
    "peak_memory_gb",
    "memory_per_generated_token_mb",
]
for col in numeric_cols:
    if col in df.columns:
        df[col] = pd.to_numeric(df[col], errors="coerce")

df = df[df["avg_decode_latency_per_token_ms"] <= FAILED_THRESHOLD].copy()
df = df.drop_duplicates()

print(f"  Rows after cleaning: {len(df)}")

# ============================================================
# AGGREGATE SUMMARY CSV
# ============================================================

summary_df = (
    df.groupby("framework")[numeric_cols]
    .agg(["mean", "median", "std", "min", "max"])
)
summary_df.to_csv(OUTPUT_DIR / "aggregate_summary.csv")
print(f"  Saved: aggregate_summary.csv")

# ============================================================
# PREP DISPLAY DATAFRAME
# ============================================================

# Flatten multi-index columns
flat = summary_df.copy()
flat.columns = [f"{metric}_{stat}" for metric, stat in flat.columns]
flat = flat.reset_index()

metric_mapping = {
    "framework": "Framework",

    "prefill_latency_s_mean":   "Prefill Mean (s)",
    "prefill_latency_s_median": "Prefill Median (s)",
    "prefill_latency_s_std":    "Prefill Std (s)",
    "prefill_latency_s_max":    "Prefill Max (s)",

    "decode_tokens_per_sec_mean":   "Decode TPS Mean",
    "decode_tokens_per_sec_median": "Decode TPS Median",
    "decode_tokens_per_sec_std":    "Decode TPS Std",

    "avg_decode_latency_per_token_ms_mean":   "Decode Lat Mean (ms)",
    "avg_decode_latency_per_token_ms_median": "Decode Lat Median (ms)",
    "avg_decode_latency_per_token_ms_std":    "Decode Lat Std (ms)",

    "peak_memory_gb_mean": "Peak Mem Mean (GB)",
    "peak_memory_gb_std":  "Peak Mem Std (GB)",

    "memory_per_generated_token_mb_mean": "Mem / Token (MB)",
}

def format_value(col, val):
    if pd.isna(val):
        return ""
    if "tokens_per_sec" in col:
        return f"{val:.2f}"
    return f"{val:.3f}"

for col in flat.columns:
    if col != "framework":
        flat[col] = flat[col].apply(lambda x: format_value(col, x))

available_cols = [c for c in metric_mapping if c in flat.columns]
display_df = (
    flat[available_cols]
    .rename(columns=metric_mapping)
    .set_index("Framework")
    .T
)

col_order = [c for c in ["Pruna", "HQQ"] if c in display_df.columns]
display_df = display_df[col_order]

# ============================================================
# STYLED TABLE FIGURE
# ============================================================

fig, ax = plt.subplots(figsize=(9, 7.5))
ax.axis("off")

col_labels = [
    "Pruna (HQQ + torch.compile)",
    "Base (HQQ + torch.compile)",
]

table = ax.table(
    cellText=display_df.values,
    rowLabels=display_df.index,
    colLabels=col_labels,
    cellLoc="center",
    rowLoc="right",
    loc="center",
)
table.auto_set_font_size(False)
table.set_fontsize(9.5)
table.scale(1.15, 1.85)

for (row, col), cell in table.get_celld().items():
    cell.set_linewidth(0.4)
    cell.set_edgecolor("#dddddd")

    if row == 0:                          # header row
        cell.set_height(0.082)
        cell.set_text_props(weight="bold", color="white", fontsize=10)
        if col == 0:
            cell.set_facecolor(PRUNA_COLOR)
        elif col == 1:
            cell.set_facecolor(HQQ_COLOR)
        else:
            cell.set_facecolor(HEADER_BG)

    elif col == -1:                       # row labels
        cell.set_text_props(weight="bold", fontsize=9, color="#333333")
        cell.set_facecolor(ROW_LABEL_BG)
        cell.set_edgecolor("#cccccc")

    else:                                 # data cells
        data_row = row - 1
        if data_row % 2 == 0:
            bg = PRUNA_LIGHT if col == 0 else HQQ_LIGHT
        else:
            bg = "white"
        cell.set_facecolor(bg)
        cell.set_text_props(
            color=PRUNA_COLOR if col == 0 else HQQ_COLOR,
            fontsize=9.5,
        )

plt.title(
    "Aggregate Benchmark Summary\nPruna  vs  HQQ + torch.compile",
    fontsize=15,
    weight="bold",
    pad=20,
    color="#1A1A2E",
)

plt.figtext(
    0.5, 0.015,
    "Values are mean / median / std across all benchmark runs "
    "after removing failed HQQ runs (latency > 1000 ms).",
    ha="center",
    fontsize=8.5,
    color="#666666",
)

output_path = OUTPUT_DIR / "aggregate_summary_table_styled.png"
plt.savefig(output_path, dpi=300, bbox_inches="tight")
plt.close()

print(f"  Saved: aggregate_summary_table_styled.png")
print(f"\nDone. Outputs in: {OUTPUT_DIR}")