EDA Script for creation of visual plots.

dataset_eda.py

"""
EDA Data Pipeline — Pruna vs HQQ
==================================
Loads raw benchmark datasets from HuggingFace Hub,
cleans them, and saves processed CSVs.

Outputs
-------
benchmark/eda_outputs/
├── combined_cleaned_results.parquet
├── aggregate_summary.csv
├── correlation_matrix.csv
└── stability_analysis.csv
"""

from pathlib import Path

import pandas as pd
from datasets import load_dataset


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

HQQ_DATASET   = "AINovice2005/raw_benchmark_results_hqq"
PRUNA_DATASET = "AINovice2005/raw_benchmark_results_pruna"

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


# ============================================================
# LOAD
# ============================================================

print("Loading datasets from HuggingFace Hub...")


def load_benchmark(repo_id: str, framework: str) -> pd.DataFrame:
    df = load_dataset(repo_id, split="train", cache_dir=CACHE_DIR).to_pandas()
    df["framework"] = framework
    return df


hqq_df   = load_benchmark(HQQ_DATASET,   "HQQ")
pruna_df = load_benchmark(PRUNA_DATASET, "Pruna")

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

print(f"  HQQ rows   : {len(hqq_df)}")
print(f"  Pruna rows : {len(pruna_df)}")
print(f"  Total rows : {len(df)}")


# ============================================================
# CLEAN
# ============================================================

numeric_cols = [
    "prompt_target_length",
    "actual_prompt_length",
    "generation_length",
    "prefill_latency_s",
    "prefill_peak_memory_gb",
    "decode_time_s",
    "decode_tokens_per_sec",
    "avg_decode_latency_per_token_ms",
    "peak_memory_gb",
    "decode_memory_growth_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.drop_duplicates()


# ============================================================
# SAVE OUTPUTS
# ============================================================

# 1. Cleaned parquet
df.to_parquet(OUTPUT_DIR / "combined_cleaned_results.parquet", index=False)
print("\n  Saved: combined_cleaned_results.parquet")

# 2. Aggregate summary
summary_metrics = [
    "prefill_latency_s",
    "decode_tokens_per_sec",
    "avg_decode_latency_per_token_ms",
    "peak_memory_gb",
    "memory_per_generated_token_mb",
]
(
    df.groupby("framework")[summary_metrics]
    .agg(["mean", "median", "std", "min", "max"])
    .to_csv(OUTPUT_DIR / "aggregate_summary.csv")
)
print("  Saved: aggregate_summary.csv")

# 3. Correlation matrix
corr_cols = [c for c in [
    "prefill_latency_s",
    "decode_tokens_per_sec",
    "avg_decode_latency_per_token_ms",
    "peak_memory_gb",
    "decode_memory_growth_gb",
    "memory_per_generated_token_mb",
    "generation_length",
] if c in df.columns]

df[corr_cols].corr().to_csv(OUTPUT_DIR / "correlation_matrix.csv")
print("  Saved: correlation_matrix.csv")

# 4. Stability analysis
stability = df.groupby("framework")["decode_tokens_per_sec"].agg(["mean", "std"])
stability["cv"] = stability["std"] / stability["mean"]
stability.to_csv(OUTPUT_DIR / "stability_analysis.csv")
print("  Saved: stability_analysis.csv")

print(f"\nDone. All outputs in: {OUTPUT_DIR}")