Update docs, timelines bug fix

Author: jaywonchung

README.md

@@ -2,32 +2,73 @@
 
 A Python toolkit for [ML.ENERGY](https://ml.energy) datasets: loading raw results, filtering and analyzing runs, fitting models, and building data packages.
 
-Actual data are stored in Hugging Face Hub: [ml-energy/benchmark-v3](https://huggingface.co/datasets/ml-energy/benchmark-v3). This repository contains code for working with the data, not the data itself.
+We currently have [The ML.ENERGY Benchmark v3.0](https://github.com/ml-energy/benchmark) dataset, which includes LLM and diffusion inference runs on NVIDIA H100 and B200 GPUs.
+Actual data are stored in Hugging Face Hub: [`ml-energy/benchmark-v3`](https://huggingface.co/datasets/ml-energy/benchmark-v3).
+This repository contains the toolkit code, not the data itself.
+
+## What the toolkit does
+
+- **Load and filter benchmark runs** with typed, immutable collection classes (`LLMRuns`, `DiffusionRuns`).
+- **Extract bulk data** — power timelines, ITL samples, output lengths — as DataFrames.
+- **Fit models** — logistic power/latency curves, ITL latency distributions.
+- **Build data packages** for publishing to Hugging Face Hub.
 
 ## Installation
 
 ```bash
-pip install -e .
+pip install mlenergy-data
 ```
 
 ## Quick example
 
 ```python
 from mlenergy_data.records import LLMRuns
 
-runs = LLMRuns.from_directory("/path/to/compiled/data")
+runs = LLMRuns.from_hf()
 
 # Find the most energy-efficient model on GPQA
 best = min(runs.task("gpqa"), key=lambda r: r.energy_per_token_joules)
-print(f"{best.nickname}: {best.energy_per_token_joules:.3f} J/tok")
+print(f"{best.nickname}: {best.energy_per_token_joules:.3f} J/tok on {best.gpu_model}")
 
 # Column access via .data
 energies = runs.data.energy_per_token_joules  # list[float]
 ```
 
+Filter, group, and compare across GPU generations and model architectures:
+
+```python
+# Compare GPU generations: best energy efficiency per model on GPQA
+for gpu, group in runs.task("gpqa").group_by("gpu_model").items():
+    best = min(group, key=lambda r: r.energy_per_token_joules)
+    print(f"{gpu}: {best.nickname} @ {best.energy_per_token_joules:.3f} J/tok, "
+          f"{best.output_throughput_tokens_per_sec:.0f} tok/s")
+
+# MoE, Dense, Hybrid: who's more energy-efficient?
+for arch, group in runs.task("gpqa").gpu("B200").group_by("architecture").items():
+    best = min(group, key=lambda r: r.energy_per_token_joules)
+    print(f"{arch}: {best.nickname} @ {best.energy_per_token_joules:.3f} J/tok")
+```
+
+## Who uses it
+
+- [**The ML.ENERGY Leaderboard v3.0**](https://ml.energy/leaderboard): Benchmark results are loaded and compiled into the leaderboard web app data format.
+- [**OpenG2G**](TODO): Datacenter-grid coordination simulation framework; loads benchmark data and fits models.
+- [**The ML.ENERGY blog**](https://ml.energy/blog): Analysis scripts for blog posts.
+
 ## Documentation
 
 See the full [documentation site](https://ml-energy.github.io/mlenergy-data/) for:
 
 - [Usage guide](https://ml-energy.github.io/mlenergy-data/guide/) — progressive walkthrough from loading data to fitting models.
 - [API reference](https://ml-energy.github.io/mlenergy-data/api/records/) — auto-generated from docstrings.
+
+## Citation
+
+```bibtex
+@inproceedings{mlenergy-neuripsdb25,
+    title={The {ML.ENERGY Benchmark}: Toward Automated Inference Energy Measurement and Optimization},
+    author={Jae-Won Chung and Jeff J. Ma and Ruofan Wu and Jiachen Liu and Oh Jun Kweon and Yuxuan Xia and Zhiyu Wu and Mosharaf Chowdhury},
+    year={2025},
+    booktitle={NeurIPS Datasets and Benchmarks},
+}
+```

data_publishing/DATASET_CARD.md

@@ -1,6 +1,6 @@
 # The ML.ENERGY Benchmark V3 Dataset
 
-This dataset contains benchmark results from [The ML.ENERGY Benchmark](https://github.com/ml-energy/benchmark).
+This dataset contains benchmark results from [The ML.ENERGY Benchmark](https://github.com/ml-energy/benchmark), which includes LLM and diffusion inference runs on NVIDIA H100 and B200 GPUs.
 You can use [The ML.ENERGY Leaderboard](https://ml.energy) to explore the benchmarking results.
 
 ## Subsets
@@ -10,7 +10,7 @@ You can use [The ML.ENERGY Leaderboard](https://ml.energy) to explore the benchm
 
 ## Usage
 
-You can programmatically utilize the dataset using the ML.ENERGY data toolkit.
+You can programmatically work with the dataset using the [ML.ENERGY data toolkit](https://github.com/ml-energy/data).
 
 ```bash
 pip install mlenergy-data
@@ -20,14 +20,36 @@ pip install mlenergy-data
 from mlenergy_data.records import LLMRuns, DiffusionRuns
 
 # Load (fast, parquet only ~few MB)
-llm = LLMRuns.from_hf()
+runs = LLMRuns.from_hf()
 
-# Filter and analyze (parquet only, no download)
-for r in llm.task("gpqa").gpu("B200"):
-    print(r.nickname, r.energy_per_token_joules)
+# Find the most energy-efficient model on GPQA
+best = min(runs.task("gpqa"), key=lambda r: r.energy_per_token_joules)
+print(f"{best.nickname}: {best.energy_per_token_joules:.3f} J/tok on {best.gpu_model}")
 
-# Bulk data methods auto-download raw files as needed
-out = llm.task("gpqa").output_lengths()
+# Column access via .data
+energies = runs.data.energy_per_token_joules  # list[float]
+```
+
+Filter, group, and compare across GPU generations and model architectures:
+
+```python
+# Compare GPU generations: best energy efficiency per model on GPQA
+for gpu, group in runs.task("gpqa").group_by("gpu_model").items():
+    best = min(group, key=lambda r: r.energy_per_token_joules)
+    print(f"{gpu}: {best.nickname} @ {best.energy_per_token_joules:.3f} J/tok, "
+          f"{best.output_throughput_tokens_per_sec:.0f} tok/s")
+
+# MoE, Dense, Hybrid: who's more energy-efficient?
+for arch, group in runs.task("gpqa").gpu("B200").group_by("architecture").items():
+    best = min(group, key=lambda r: r.energy_per_token_joules)
+    print(f"{arch}: {best.nickname} @ {best.energy_per_token_joules:.3f} J/tok")
+```
+
+Bulk data methods auto-download raw files as needed:
+
+```python
+# Power timelines, output lengths, etc.
+power_tl = runs.timelines(metric="power.device_instant")
 ```
 
 ## Schema
@@ -74,7 +96,7 @@ Please direct any issues, questions, or discussions to [The ML.ENERGY Data Toolk
 
 ```bibtex
 @inproceedings{mlenergy-neuripsdb25,
-    title={The {ML.ENERGY Benchmark}: Toward Automated Inference Energy Measurement and Optimization}, 
+    title={The {ML.ENERGY Benchmark}: Toward Automated Inference Energy Measurement and Optimization},
     author={Jae-Won Chung and Jeff J. Ma and Ruofan Wu and Jiachen Liu and Oh Jun Kweon and Yuxuan Xia and Zhiyu Wu and Mosharaf Chowdhury},
     year={2025},
     booktitle={NeurIPS Datasets and Benchmarks},

docs/guide.md

@@ -4,9 +4,19 @@
 
 For full working examples of the toolkit in production, see:
 
-- [ML.ENERGY Leaderboard data build](https://github.com/ml-energy/leaderboard/blob/main/scripts/build_data.py) -- builds the leaderboard JSON data from benchmark runs
-- [ML.ENERGY Blog analysis scripts](https://github.com/ml-energy/data/blob/main/blog_analysis_scripts.py) -- generates figures for the ML.ENERGY blog
-- [OpenG2G simulation data build](TODO) -- builds power traces, logistic fits, and latency fits for grid simulation
+- ML.ENERGY Leaderboard data build
+   - Builds the leaderboard JSON data from benchmark runs
+   - [The ML.ENERGY Leaderboard](https://ml.energy/leaderboard)
+   - [Data build script](https://github.com/ml-energy/leaderboard/blob/main/scripts/build_data.py)
+- ML.ENERGY Blog analysis scripts
+   - Generates figures for the ML.ENERGY blog post on the V3 benchmark results
+   - [Blog post](https://ml.energy/blog/measurement/energy/diagnosing-inference-energy-consumption-with-the-mlenergy-leaderboard-v30/)
+   - [Analysis script](TODO)
+- OpenG2G simulation data build
+   - Builds power traces, logistic fits, and latency distributions for datacenter--grid coordination simulation
+   - [OpenG2G](TODO)
+   - [Data build script](TODO)
+
 
 ## Loading benchmark runs
 
@@ -17,7 +27,8 @@ Each run is a frozen dataclass (`LLMRun` / `DiffusionRun`) with IDE autocomplete
 from mlenergy_data.records import LLMRuns, DiffusionRuns
 
 # Load all stable LLM runs from a compiled data directory
-runs = LLMRuns.from_directory("/path/to/compiled/data")
+root = "/path/to/compiled/data"
+runs = LLMRuns.from_directory(root)
 
 # Include unstable runs
 runs = LLMRuns.from_directory(root, stable_only=False)
@@ -30,6 +41,9 @@ runs = LLMRuns.from_hf()
 diff = DiffusionRuns.from_hf()
 ```
 
+!!! Note
+    A "compiled data directory" is one built by `data_publishing/build_hf_data.py` (or downloaded from HF Hub). It contains parquet summary files under `runs/`, raw result files under `llm/` and `diffusion/`, and benchmark config files under `configs/`.
+
 ## Filtering
 
 All filter methods return a new collection — chain freely:
@@ -49,6 +63,9 @@ chat_or_gpqa = runs.task("gpqa", "lm-arena-chat")
 # By nickname
 deepseek = runs.nickname("DeepSeek R1")
 
+# Architecture (LLM only)
+moe_models = runs.architecture("MoE")
+
 # Batch size: exact values or range
 batch_128 = runs.batch(128)
 large_batch = runs.batch(min=64)
@@ -58,6 +75,11 @@ mid_batch = runs.batch(min=16, max=128)
 single_gpu = runs.num_gpus(1)
 multi_gpu = runs.num_gpus(min=2)
 
+# Stability
+# Relevant when you explicitly set stable_only=False at load time to include unstable runs. By default, only stable runs are loaded.
+stable_only = runs.stable()
+unstable_only = runs.unstable()
+
 # Arbitrary predicate
 big_models = runs.where(lambda r: r.total_params_billions > 70)
 ```
@@ -95,6 +117,17 @@ plt.xlabel("Batch size")
 plt.ylabel("Energy per token (J)")
 ```
 
+**Indexing and concatenation:**
+
+```python
+first_run = runs[0]
+
+# Concatenate collections
+h100 = runs.gpu("H100")
+b200 = runs.gpu("B200")
+combined = h100 + b200
+```
+
 ## Grouping
 
 ```python
@@ -113,12 +146,12 @@ for (model, batch), g in runs.group_by("model_id", "max_num_seqs").items():
 Python is the analysis layer — no special helper functions needed:
 
 ```python
-# Best energy per token for each model on a task
-for model_id, group in runs.task("gpqa").group_by("model_id").items():
-    best = min(group, key=lambda r: r.energy_per_token_joules)
-    print(f"{best.nickname}: {best.energy_per_token_joules:.3f} J/tok")
+# Compare GPU generations on a task
+for gpu, group in runs.task("lm-arena-chat").group_by("gpu_model").items():
+    best = min(group, key=lambda r: r.output_throughput_tokens_per_sec)
+    print(f"{gpu}: {best.nickname} @ {best.output_throughput_tokens_per_sec:.0f} tok/s")
 
-# Comparing GPUs
+# Comparing GPUs for a specific model
 llama70b = runs.model("meta-llama/Llama-3.1-70B-Instruct")
 for gpu, g in llama70b.group_by("gpu_model").items():
     plt.scatter(g.data.max_num_seqs, g.data.energy_per_token_joules, label=gpu)
@@ -130,6 +163,14 @@ plt.legend()
 These methods return pandas DataFrames for numerical analysis.
 When loaded from HF Hub (`from_hf()`), they automatically download only the raw files needed for the current collection. The download scope is determined by your filters. HF Hub caches files locally, so repeated calls are instant.
 
+To eagerly download all raw files upfront, use `prefetch()`:
+
+```python
+# Eagerly download all raw files for a filtered collection
+runs = LLMRuns.from_hf().task("gpqa").prefetch()
+power_tl = runs.timelines(metric="power.device_instant")  # no download delay
+```
+
 ```python
 # Power timelines (long-form)
 power_tl = runs.timelines(metric="power.device_instant")
@@ -158,6 +199,12 @@ t2i = diff.task("text-to-image")
 best = min(t2i, key=lambda r: r.energy_per_generation_joules)
 print(f"{best.nickname}: {best.energy_per_generation_joules:.3f} J/image")
 
+# Task field and convenience properties
+r = diff[0]
+r.task               # "text-to-image" or "text-to-video"
+r.is_text_to_image   # True for text-to-image tasks
+r.is_text_to_video   # True for text-to-video tasks
+
 # Available filters: task(), model(), gpu(), nickname(), batch(),
 # num_gpus(), precision(), where()
 ```
@@ -206,15 +253,3 @@ avg_itl = model.sample_avg(n_replicas=180, rng=rng)
 d = model.to_dict()
 model2 = ITLMixtureModel.from_dict(d)
 ```
-
-## Building a Hugging Face data package
-
-```bash
-python data_publishing/build_hf_data.py \
-  --results-dir /path/to/llm/h100/current/run \
-  --results-dir /path/to/diffusion/h100/current/run \
-  --out-dir /tmp/hf_pkg
-
-# Upload to Hugging Face Hub
-hf upload-large-folder ml-energy/benchmark-v3 /tmp/hf_pkg --repo-type dataset
-```

docs/index.md

@@ -4,7 +4,7 @@
 To aid in working with these datasets, we also provide a Python toolkit: `mlenergy-data`.
 
 We currently have [The ML.ENERGY Benchmark v3.0](https://github.com/ml-energy/benchmark) dataset, which includes LLM and diffusion inference runs on NVIDIA H100 and B200 GPUs.
-Actual data are currently stored in Hugging Face Hub: [`ml-energy/benchmark-v3`](https://huggingface.co/datasets/ml-energy/benchmark-v3).
+Actual data are stored in Hugging Face Hub: [`ml-energy/benchmark-v3`](https://huggingface.co/datasets/ml-energy/benchmark-v3).
 
 ## What the toolkit does
 
@@ -29,6 +29,24 @@ runs = LLMRuns.from_hf()
 # Find the most energy-efficient model on GPQA
 best = min(runs.task("gpqa"), key=lambda r: r.energy_per_token_joules)
 print(f"{best.nickname}: {best.energy_per_token_joules:.3f} J/tok on {best.gpu_model}")
+
+# Column access via .data
+energies = runs.data.energy_per_token_joules  # list[float]
+```
+
+Filter, group, and compare across GPU generations and model architectures:
+
+```python
+# Compare GPU generations: best energy efficiency per model on GPQA
+for gpu, group in runs.task("gpqa").group_by("gpu_model").items():
+    best = min(group, key=lambda r: r.energy_per_token_joules)
+    print(f"{gpu}: {best.nickname} @ {best.energy_per_token_joules:.3f} J/tok, "
+          f"{best.output_throughput_tokens_per_sec:.0f} tok/s")
+
+# MoE, Dense, Hybrid: who's more energy-efficient?
+for arch, group in runs.task("gpqa").gpu("B200").group_by("architecture").items():
+    best = min(group, key=lambda r: r.energy_per_token_joules)
+    print(f"{arch}: {best.nickname} @ {best.energy_per_token_joules:.3f} J/tok")
 ```
 
 ## Who uses it
@@ -41,3 +59,14 @@ print(f"{best.nickname}: {best.energy_per_token_joules:.3f} J/tok on {best.gpu_m
 
 - [Guide](guide.md): Progressive walkthrough from loading data to fitting models.
 - [API Reference](api/records.md): Auto-generated from docstrings.
+
+## Citation
+
+```bibtex
+@inproceedings{mlenergy-neuripsdb25,
+    title={The {ML.ENERGY Benchmark}: Toward Automated Inference Energy Measurement and Optimization},
+    author={Jae-Won Chung and Jeff J. Ma and Ruofan Wu and Jiachen Liu and Oh Jun Kweon and Yuxuan Xia and Zhiyu Wu and Mosharaf Chowdhury},
+    year={2025},
+    booktitle={NeurIPS Datasets and Benchmarks},
+}
+```

mlenergy_data/modeling/latency.py

@@ -38,7 +38,11 @@ def _lognormal_mean_var(self, sigma: float, scale: float) -> tuple[float, float]
         return ey, vy
 
     def mean_var(self) -> tuple[float, float]:
-        """Compute mean and variance of the two-component mixture."""
+        """Compute analytical mean and variance of the two-component mixture.
+
+        Returns:
+            Tuple of (mean, variance) in seconds.
+        """
         p1 = float(self.pi_steady)
         p2 = float(self.pi_stall)
         ps = max(p1 + p2, 1e-12)
@@ -56,7 +60,14 @@ def mean_var(self) -> tuple[float, float]:
         return mx, vx
 
     def sample_one(self, rng: np.random.Generator) -> float:
-        """Draw a single sample from the mixture."""
+        """Draw a single ITL sample from the mixture.
+
+        Args:
+            rng: NumPy random generator.
+
+        Returns:
+            Sampled ITL value in seconds.
+        """
         p1 = float(self.pi_steady)
         p2 = float(self.pi_stall)
         ps = max(p1 + p2, 1e-12)

mlenergy_data/modeling/logistic.py

@@ -28,7 +28,7 @@ class LogisticModel:
     b0: float
 
     def eval_x(self, x: float) -> float:
-        """Evaluate at continuous x (= log2(batch_size))."""
+        """Evaluate at continuous x (typically log2(batch_size))."""
         a = self.k * (float(x) - self.x0)
         if a >= 0:
             ea = math.exp(-a)
@@ -39,7 +39,7 @@ def eval_x(self, x: float) -> float:
         return float(self.b0 + self.L * s)
 
     def deriv_wrt_x(self, x: float) -> float:
-        """dy/dx for y = b0 + L * sigmoid(k*(x - x0))."""
+        """Derivative dy/dx at continuous x."""
         a = self.k * (float(x) - self.x0)
         if a >= 0:
             ea = math.exp(-a)