Spaces:
Sleeping
Sleeping
File size: 3,141 Bytes
24e1932 a81fd74 24e1932 a07e79a 24e1932 | 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 | """
Tests for the transform stage (reprocess_metrics logic).
These tests exercise the metric re-computation on an existing on-disk dataset
without downloading any model weights from HuggingFace.
"""
import os
import numpy as np
import pandas as pd
import pytest
from datasets import load_from_disk
from transformer_analysis.head_metrics import normality_metrics, singular_value_metrics
NORMALITY_KEYS = list(normality_metrics.keys())
SV_KEYS = list(singular_value_metrics.keys())
def _recompute_on_dataset(dataset_dir):
"""
Core of reprocess_metrics logic, extracted for direct testing without
hitting the model_name / out_dir path resolution in the full function.
"""
import json
import numpy as np
ds = load_from_disk(str(dataset_dir))
df = ds.to_pandas()
metadata_path = os.path.join(str(dataset_dir), "metadata.json")
with open(metadata_path) as f:
meta = json.load(f)
w_bins = np.array(meta["w_bins"])
centers = (w_bins[:-1] + w_bins[1:]) / 2
new_columns = {}
for idx, row in df.iterrows():
h = row.to_dict()
p_w = h.get("P_w")
if p_w is not None:
try:
p_w = np.asarray(p_w, dtype=float)
if not np.any(np.isnan(p_w)):
h["P_w"] = p_w
for fn in normality_metrics.values():
fn(h, centers)
except (TypeError, ValueError):
pass
svd = h.get("SVD")
if svd is not None and hasattr(svd, "__len__") and len(svd) > 0:
for fn in singular_value_metrics.values():
fn(h, svd)
for key, value in h.items():
if key not in new_columns:
new_columns[key] = [None] * len(df)
new_columns[key][idx] = value
for col, vals in new_columns.items():
df[col] = vals
return df
def test_transform_adds_normality_columns(tiny_dataset_dir):
df = _recompute_on_dataset(tiny_dataset_dir)
for key in ["entropy", "fit_mu", "fit_sigma", "kl_vs_empirical_normal"]:
assert key in df.columns, f"Missing column: {key}"
def test_transform_adds_sv_columns(tiny_dataset_dir):
df = _recompute_on_dataset(tiny_dataset_dir)
wqk_rows = df[df["weight_type"] == "W_QK"]
assert len(wqk_rows) > 0
for key in ["normalized_participation_ratio", "spectral_entropy", "condition_number"]:
assert key in df.columns, f"Missing column: {key}"
assert wqk_rows[key].notna().all(), f"NaN values in {key} for W_QK rows"
def test_transform_preserves_row_count(tiny_dataset_dir):
ds_before = load_from_disk(str(tiny_dataset_dir))
df_after = _recompute_on_dataset(tiny_dataset_dir)
assert len(df_after) == len(ds_before)
def test_transform_reproducible(tiny_dataset_dir):
df1 = _recompute_on_dataset(tiny_dataset_dir)
df2 = _recompute_on_dataset(tiny_dataset_dir)
for col in ["entropy", "normalized_participation_ratio", "spectral_entropy"]:
assert np.allclose(df1[col].values, df2[col].values, equal_nan=True), \
f"Column {col} not reproducible across runs"
|