bdrgon's picture
Validate Gemma exports on full drum corpus and publish results
448ab1d
Raw History Blame Contribute Delete
25.4 kB
"""Full-corpus, model-only validation shared by the independent exporters.
No native application/library is required. Every dataset row is accounted for;
bad inputs and inference failures are retained in the report, never silently dropped.
"""
from __future__ import annotations
import argparse
from concurrent.futures import ProcessPoolExecutor
from collections import defaultdict
from datetime import datetime, timezone
import gc
import hashlib
import importlib.metadata
import io
import json
import math
from pathlib import Path
import platform
import time
import numpy as np
from export_common import MODEL_ID, REVISION, ROOT, load_reference, sha256_file
DATASET = "yojul/one-shot-hip-hop-drums"
DATASET_REVISION = "07be604e729fcf8f7904bca5328d6b09b1ceed7b"
DATASET_ROWS = 19673
VARIANTS = ("fp32", "int8", "int4")
COMPARISONS = ("upstream -> fp32", "fp32 -> int8", "fp32 -> int4")
def parser(description):
result = argparse.ArgumentParser(description=description)
mode = result.add_mutually_exclusive_group()
mode.add_argument("--validate-only", action="store_true", help="Validate existing graphs without exporting again")
mode.add_argument("--skip-validation", action="store_true", help="Export with smoke checks only; skip corpus evaluation")
result.add_argument("--dataset-cache", type=Path, help="Hugging Face cache (default: standard HF cache)")
result.add_argument("--results-dir", type=Path, default=ROOT / "validation-results")
result.add_argument("--work-dir", type=Path, default=ROOT / ".validation-cache")
result.add_argument("--threads", type=int, default=2)
result.add_argument("--workers", type=int, default=1, help="Independent CPU workers; each loads its own models")
result.add_argument("--limit", type=int, help="Diagnostic prefix only; visibly marked PARTIAL, never a full-corpus result")
result.add_argument("--resume", action="store_true", help="Resume only an exactly matching checkpoint")
result.add_argument("--offline", action="store_true", help="Use cached Hub files only")
return result
def divergence(reference, actual):
if reference.shape != actual.shape:
raise ValueError(f"shape mismatch: {reference.shape} != {actual.shape}")
a, b = reference.astype(np.float64).ravel(), actual.astype(np.float64).ravel()
if not len(a) or not np.isfinite(a).all() or not np.isfinite(b).all():
raise ValueError("empty or non-finite output")
denominator = float(np.linalg.norm(a) * np.linalg.norm(b))
if denominator <= 0:
raise ValueError("zero-norm output has undefined cosine")
difference = a - b
return {
"cosine": float(np.clip(np.dot(a, b) / denominator, -1, 1)),
"max_abs": float(np.max(np.abs(difference))),
"rmse": float(np.sqrt(np.mean(difference * difference))),
}
def summary(values):
values = np.asarray(values, dtype=np.float64)
if not len(values):
return {key: None for key in ("mean", "min", "p05", "p95", "max")}
return {"mean": float(values.mean()), "min": float(values.min()),
"p05": float(np.percentile(values, 5)), "p95": float(np.percentile(values, 95)),
"max": float(values.max())}
def passes(metrics, comparison):
# Frozen exporter tolerances. Quantized values are observations, not a claim
# that these permissive cosine thresholds imply perceptual quality.
if comparison == "upstream -> fp32":
return metrics["cosine"] >= 0.99999 and metrics["max_abs"] <= 5e-4
return metrics["cosine"] >= (0.97 if comparison.endswith("int8") else 0.80)
def decode_audio(row):
import soundfile as sf
from scipy.signal import resample_poly
encoded = row["audio"]["bytes"]
if not encoded:
raise ValueError("dataset row contains zero audio bytes")
samples, rate = sf.read(io.BytesIO(encoded), dtype="float32", always_2d=True)
source_frames, channels = samples.shape
if not source_frames or not np.isfinite(samples).all():
raise ValueError("empty or non-finite decoded audio")
# Explicit shared Python preparation, independent of any native runtime.
mono = samples[:30 * rate].mean(axis=1, dtype=np.float32)
if rate != 16000:
divisor = math.gcd(rate, 16000)
mono = resample_poly(mono, 16000 // divisor, rate // divisor).astype(np.float32)
if len(mono) <= 160:
raise ValueError("Gemma requires more than 10 ms of audio")
return mono, {"sample_rate": rate, "channels": channels, "source_frames": source_frames,
"evaluated_samples": len(mono), "truncated": source_frames > 30 * rate}
def make_inputs(processor, samples):
inputs = processor(audio=[samples], sampling_rate=16000, return_tensors="np")
dtypes = {"input_ids": np.int64, "attention_mask": np.int64,
"input_features": np.float32, "input_features_mask": np.bool_}
return {name: inputs[name].astype(dtype) for name, dtype in dtypes.items()}
def load_upstream(kind, config):
import torch
if kind == "tower":
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from transformers.models.gemma4.modeling_gemma4 import Gemma4AudioModel
from export_tower import AudioTower
model = Gemma4AudioModel(config.audio_config)
checkpoint = hf_hub_download(MODEL_ID, "model.safetensors", revision=REVISION)
with safe_open(checkpoint, framework="pt", device="cpu") as weights:
state = {key.removeprefix("audio_tower."): weights.get_tensor(key)
for key in weights.keys() if key.startswith("audio_tower.")}
model.load_state_dict(state, strict=True)
return AudioTower(model).eval()
from transformers import AutoModel
from export_embedding import AudioEmbedding
config.vision_config = None
return AudioEmbedding(AutoModel.from_pretrained(
MODEL_ID, revision=REVISION, config=config, dtype=torch.float32,
).eval()).eval()
def upstream_output(model, kind, feed):
import torch
tensors = {name: torch.from_numpy(value) for name, value in feed.items()}
with torch.inference_mode():
if kind == "tower":
states, mask = model(tensors["input_features"], tensors["input_features_mask"])
return states.cpu().float().numpy(), mask.cpu().numpy()
return model(**tensors).cpu().float().numpy(), None
def ort_output(session, feed, kind):
outputs = session.run(None, {value.name: feed[value.name] for value in session.get_inputs()})
return outputs[0], outputs[1] if kind == "tower" else None
def output_contract(values, mask, kind):
if not np.isfinite(values).all():
raise ValueError("non-finite model output")
if kind == "embedding":
if values.shape != (1, 512) or not np.isclose(np.linalg.norm(values), 1, atol=1e-5):
raise ValueError("embedding shape or L2 normalization contract failed")
elif (values.ndim != 3 or values.shape[0] != 1 or values.shape[2] != 1536
or mask is None or mask.shape != values.shape[:2] or mask.dtype != np.bool_):
raise ValueError("tower shape/mask contract failed")
def bucket_key(shape, mask):
# Raw towers are never pooled or compared across incompatible shapes/masks.
return json.dumps([shape, mask], separators=(",", ":"))
def nearest_indices(vectors, k):
"""Exact cosine neighbors, self excluded; row index breaks exact ties."""
norms = np.linalg.norm(vectors.astype(np.float64), axis=1)
normalized = (vectors / norms[:, None]).astype(np.float32)
output = np.empty((len(vectors), k), dtype=np.int32)
for start in range(0, len(vectors), 128):
scores = normalized[start:start + 128] @ normalized.T
scores[np.arange(len(scores)), np.arange(start, start + len(scores))] = -np.inf
output[start:start + len(scores)] = np.argsort(-scores, axis=1, kind="stable")[:, :k]
return output
def neighbor_agreement(records, work):
groups = defaultdict(list)
for record in records:
if record["status"] == "ok":
groups[bucket_key(record["shape"], record["mask"])].append(record)
maps = {variant: np.memmap(work / f"{variant}.f32", dtype="<f4", mode="r")
for variant in VARIANTS} if any(groups.values()) else {}
totals = {variant: {"eligible": 0, "top1_matches": 0, "overlap_at_10": [], "top10_eligible": 0}
for variant in ("int8", "int4")}
singleton = 0
for group_number, group in enumerate(groups.values()):
if len(group) < 2:
singleton += len(group)
continue
k = min(10, len(group) - 1)
def matrix(variant):
return np.stack([maps[variant][r["offset"]:r["offset"] + r["count"]] for r in group])
reference = nearest_indices(matrix("fp32"), k)
for variant in ("int8", "int4"):
actual = nearest_indices(matrix(variant), k)
total = totals[variant]
total["eligible"] += len(group)
total["top1_matches"] += int((reference[:, 0] == actual[:, 0]).sum())
if k == 10:
total["top10_eligible"] += len(group)
total["overlap_at_10"].extend(
len(set(a) & set(b)) / 10 for a, b in zip(reference, actual)
)
if group_number % 50 == 0:
print(f"Neighbor groups: {group_number + 1}/{len(groups)}", flush=True)
return {variant: {"eligible": total["eligible"], "singleton_excluded": singleton,
"top1_agreement": total["top1_matches"] / total["eligible"] if total["eligible"] else None,
"top10_eligible": total["top10_eligible"],
"mean_overlap_at_10": float(np.mean(total["overlap_at_10"])) if total["overlap_at_10"] else None}
for variant, total in totals.items()}
SWEEPS = [("pitch", [110., 220., 440.]), ("pitch", [220., 330., 440.]),
("harmonic_mix", [0., .3, .9]), ("harmonic_mix", [.1, .5, 1.]),
("filter_cutoff", [300., 1000., 5000.]), ("filter_cutoff", [500., 2000., 8000.]),
("decay", [.05, .2, .8]), ("decay", [.1, .4, 1.]),
("fm_depth", [0., 1., 4.]), ("fm_depth", [.5, 2., 6.])]
def synth_signal(knob, value):
"""Frozen one-second, one-knob sweeps implemented entirely in Python."""
rate = 16000
t = np.arange(rate, dtype=np.float64) / rate
hz = value if knob == "pitch" else 220.
phase = 2 * np.pi * hz * t
fm = value * np.sin(2 * np.pi * 110 * t) if knob == "fm_depth" else 0.
harmonic = value if knob == "harmonic_mix" else .3
raw = np.sin(phase + fm) + harmonic * np.sin(2 * phase) + harmonic * .5 * np.sin(3 * phase)
cutoff = value if knob == "filter_cutoff" else 8000.
alpha = 1 - math.exp(-2 * math.pi * cutoff / rate)
low = 0.
filtered = np.empty(rate)
for i, sample in enumerate(raw):
low += alpha * (sample - low)
filtered[i] = low
decay = value if knob == "decay" else .6
envelope = np.minimum(t / .01, 1) * np.exp(-t / decay) * np.clip(1 - t, 0, .01) / .01
return (.3 * filtered * envelope).astype(np.float32)
def continuity(model, sessions, processor, kind):
results = {variant: [] for variant in ("upstream", *VARIANTS)}
for knob, values in SWEEPS:
outputs = {variant: [] for variant in results}
for value in values:
feed = make_inputs(processor, synth_signal(knob, value))
outputs["upstream"].append(upstream_output(model, kind, feed)[0])
for variant, session in sessions.items():
outputs[variant].append(ort_output(session, feed, kind)[0])
for variant, (a, b, c) in outputs.items():
distances = [1 - divergence(a, b)["cosine"], 1 - divergence(b, c)["cosine"],
1 - divergence(a, c)["cosine"]]
results[variant].append({"knob": knob, "values": values, "distances": distances,
"passed": max(distances[:2]) <= distances[2] + 1e-6})
return {variant: {"passed": sum(case["passed"] for case in cases), "total": len(cases),
"required_fraction": .9, "expectation_met": sum(c["passed"] for c in cases) >= 9,
"cases": cases} for variant, cases in results.items()}
def markdown(report):
def number(value):
return "N/A" if value is None else f"{value:.6g}"
lines = [f"### {report['kind'].title()} — {'FULL CORPUS' if report['full_corpus'] else 'PARTIAL'}",
"", f"Dataset `{DATASET}` at `{DATASET_REVISION}`.",
f"Attempted **{report['attempted']:,}/{report['dataset_rows']:,}** rows; "
f"successful **{report['successful']:,}**; input errors **{report['input_errors']:,}**; "
f"inference errors **{report['inference_errors']:,}**.", "",
"| Comparison | Clips | Within tolerance | Mean cosine | Min cosine | P95 clip max abs | Worst abs | P95 RMSE |",
"| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"]
for name, row in report["divergence"].items():
lines.append(f"| {name} | {row['count']:,} | {row['passed']:,}/{row['count']:,} | "
f"{number(row['cosine']['mean'])} | {number(row['cosine']['min'])} | "
f"{number(row['max_abs']['p95'])} | {number(row['max_abs']['max'])} | {number(row['rmse']['p95'])} |")
lines += ["", "| Quantization vs FP32 | Top-1 eligible queries | Top-1 agreement | Top-10 eligible queries | Mean overlap@10 |",
"| --- | ---: | ---: | ---: | ---: |"]
for variant, row in report["neighbors"].items():
lines.append(f"| {variant.upper()} | {row['eligible']:,} | {number(row['top1_agreement'])} | "
f"{row['top10_eligible']:,} | {number(row['mean_overlap_at_10'])} |")
lines += ["", "| Synth continuity | Passing triples | Meets 90% expectation |", "| --- | ---: | --- |"]
for variant, row in report["continuity"].items():
lines.append(f"| {variant} | {row['passed']}/{row['total']} | {'Yes' if row['expectation_met'] else 'No'} |")
lines += ["", "Raw-tower neighbor searches use only identical shapes and masks; no pooling is applied. "
"Final embeddings search all successful clips. Self matches are excluded; exact ties use dataset row order. "
"Overlap@10 includes only groups with at least 11 clips.",
"", "FP32 tolerance: cosine >= 0.99999 and max absolute error <= 0.0005. "
"INT8/INT4 cosine tolerances are 0.97/0.80, matching the existing smoke gates; "
"they are numerical diagnostics, not perceptual acceptance criteria.",
"", "Python-versus-native preprocessing parity is **deferred** and was not measured here.", ""]
return "\n".join(lines)
_CONTEXT = None
def initialize_worker(kind, threads, offline):
"""Process-local pristine reference and ONNX sessions; no export patches."""
import onnxruntime as ort
import torch
if offline:
import os
import huggingface_hub.constants
os.environ["HF_HUB_OFFLINE"] = "1"
huggingface_hub.constants.HF_HUB_OFFLINE = True
graph_names = {v: f"audio_{kind}{'' if v == 'fp32' else '_' + v}.onnx" for v in VARIANTS}
config, processor = load_reference()
torch.set_num_threads(threads)
upstream = load_upstream(kind, config)
sessions = {}
for variant, name in graph_names.items():
options = ort.SessionOptions()
options.intra_op_num_threads = threads
options.inter_op_num_threads = 1
options.log_severity_level = 3
sessions[variant] = ort.InferenceSession(str(ROOT / name), options, providers=["CPUExecutionProvider"])
global _CONTEXT
_CONTEXT = kind, processor, upstream, sessions
def evaluate_case(item):
row_number, row = item
kind, processor, upstream, sessions = _CONTEXT
record = {"row": row_number, "path": row["audio"]["path"], "label": row["label"]}
try:
samples, preparation = decode_audio(row)
feed = make_inputs(processor, samples)
record.update(preparation)
except Exception as error:
record.update(status="input_error", error=f"{type(error).__name__}: {error}")
return record, None
try:
reference, ref_mask = upstream_output(upstream, kind, feed)
output_contract(reference, ref_mask, kind)
outputs = {}
for variant, session in sessions.items():
values, mask = ort_output(session, feed, kind)
output_contract(values, mask, kind)
if not np.array_equal(mask, ref_mask):
raise ValueError(f"{variant} mask differs from upstream")
outputs[variant] = values
metrics = {"upstream -> fp32": divergence(reference, outputs["fp32"]),
"fp32 -> int8": divergence(outputs["fp32"], outputs["int8"]),
"fp32 -> int4": divergence(outputs["fp32"], outputs["int4"])}
record.update(status="ok", metrics=metrics, count=outputs["fp32"].size,
shape=list(reference.shape), mask=ref_mask.ravel().tolist() if ref_mask is not None else [])
return record, outputs
except Exception as error:
record.update(status="inference_error", error=f"{type(error).__name__}: {error}")
return record, None
def run_validation(kind, args):
global _CONTEXT
import onnxruntime as ort
import pyarrow.parquet as pq
import torch
from huggingface_hub import snapshot_download
if args.threads < 1 or args.workers < 1 or (args.limit is not None and args.limit < 1):
raise ValueError("threads and limit must be positive")
if args.offline:
# Set the library's constant too: huggingface_hub may have been imported earlier.
import os
import huggingface_hub.constants
os.environ["HF_HUB_OFFLINE"] = "1"
huggingface_hub.constants.HF_HUB_OFFLINE = True
snapshot = Path(snapshot_download(DATASET, repo_type="dataset", revision=DATASET_REVISION,
cache_dir=args.dataset_cache, local_files_only=args.offline))
shards = sorted((snapshot / "data").glob("*.parquet"))
total = sum(pq.ParquetFile(path).metadata.num_rows for path in shards)
if total != DATASET_ROWS:
raise ValueError(f"Pinned corpus row count changed: {total} != {DATASET_ROWS}")
target = min(args.limit, total) if args.limit else total
graph_names = {v: f"audio_{kind}{'' if v == 'fp32' else '_' + v}.onnx" for v in VARIANTS}
hashes = {name: sha256_file(ROOT / name) for name in graph_names.values()}
if kind == "embedding":
hashes["audio_embedding.data"] = sha256_file(ROOT / "audio_embedding.data")
versions = {name: importlib.metadata.version(name) for name in
("torch", "transformers", "onnxruntime", "numpy", "scipy", "soundfile", "pyarrow", "huggingface_hub")}
code = {name: sha256_file(ROOT / name) for name in ("validation.py", "export_common.py", f"export_{kind}.py")}
provenance = {"kind": kind, "dataset": DATASET, "dataset_revision": DATASET_REVISION,
"dataset_files": {p.name: sha256_file(p) for p in shards}, "dataset_rows": total,
"target": target, "model": MODEL_ID, "model_revision": REVISION,
"artifacts": hashes, "versions": versions, "code": code, "threads": args.threads, "workers": args.workers,
"python": platform.python_version(), "platform": platform.platform()}
identity = hashlib.sha256(json.dumps(provenance, sort_keys=True).encode()).hexdigest()
work = args.work_dir / f"{kind}-{identity[:16]}"
work.mkdir(parents=True, exist_ok=True)
log_path = work / "cases.jsonl"
if log_path.exists() and not args.resume:
raise ValueError(f"Checkpoint already exists; use --resume or another --work-dir: {work}")
records = []
if log_path.exists():
with log_path.open("rb") as stream:
good_end = 0
for line in stream:
try:
records.append(json.loads(line))
except json.JSONDecodeError:
break # A interrupted final write is not a completed row.
good_end = stream.tell()
with log_path.open("r+b") as stream:
stream.truncate(good_end)
if [r["row"] for r in records] != list(range(len(records))):
raise ValueError("Checkpoint row sequence is corrupt")
offset = sum(r.get("count", 0) for r in records if r["status"] == "ok")
streams = {}
for variant in VARIANTS:
path = work / f"{variant}.f32"
stream = path.open("r+b" if path.exists() else "w+b")
if path.stat().st_size < offset * 4:
raise ValueError("Checkpoint vectors are incomplete")
stream.truncate(offset * 4)
stream.seek(0, 2)
streams[variant] = stream
print(f"EXPERIMENTAL {kind}: {len(records):,}/{target:,} rows checkpointed; cache {work}", flush=True)
started, resumed = time.monotonic(), len(records)
def remaining_rows():
row_number = -1
for shard in shards:
for batch in pq.ParquetFile(shard).iter_batches(batch_size=128):
for row in batch.to_pylist():
row_number += 1
if row_number >= target:
return
if row_number >= resumed:
yield row_number, row
pool = None
try:
if args.workers == 1:
initialize_worker(kind, args.threads, args.offline)
results = map(evaluate_case, remaining_rows())
else:
pool = ProcessPoolExecutor(max_workers=args.workers, initializer=initialize_worker,
initargs=(kind, args.threads, args.offline))
results = pool.map(evaluate_case, remaining_rows(), chunksize=4)
with log_path.open("a", encoding="utf-8") as log:
for record, outputs in results:
if record["status"] == "ok":
record["offset"] = offset
for variant in VARIANTS:
streams[variant].write(outputs[variant].astype("<f4").tobytes())
streams[variant].flush()
offset += record["count"]
log.write(json.dumps(record, allow_nan=False) + "\n")
log.flush()
records.append(record)
completed = len(records) - resumed
if completed == 1 or len(records) % 50 == 0 or len(records) == target:
elapsed = time.monotonic() - started
speed = completed / max(elapsed, .001)
print(f"{kind}: {len(records):,}/{target:,}; {speed:.2f} rows/s; "
f"ETA {(target - len(records)) / max(speed, .001) / 60:.1f} min", flush=True)
finally:
if pool is not None:
pool.shutdown(wait=True, cancel_futures=True)
for stream in streams.values():
stream.close()
if args.workers != 1:
initialize_worker(kind, args.threads, args.offline)
_, processor, upstream, sessions = _CONTEXT
synth = continuity(upstream, sessions, processor, kind)
_CONTEXT = None
del upstream, sessions
gc.collect()
neighbors = neighbor_agreement(records, work)
successful = [r for r in records if r["status"] == "ok"]
report = dict(provenance, generated_utc=datetime.now(timezone.utc).isoformat(), full_corpus=target == total,
attempted=len(records), successful=len(successful),
input_errors=sum(r["status"] == "input_error" for r in records),
inference_errors=sum(r["status"] == "inference_error" for r in records),
truncated=sum(r.get("truncated", False) for r in records),
divergence={}, neighbors=neighbors, continuity=synth,
errors=[r for r in records if r["status"] != "ok"],
deferred=["Python-versus-native preprocessing parity"], checkpoint_identity=identity)
for comparison in COMPARISONS:
values = [r["metrics"][comparison] for r in successful]
report["divergence"][comparison] = {"count": len(values),
"passed": sum(passes(v, comparison) for v in values),
**{metric: summary([v[metric] for v in values]) for metric in ("cosine", "max_abs", "rmse")}}
args.results_dir.mkdir(parents=True, exist_ok=True)
label = kind if target == total else f"{kind}-partial-{target}"
(args.results_dir / f"{label}.json").write_text(json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8")
table = markdown(report)
(args.results_dir / f"{label}.md").write_text(table, encoding="utf-8")
print("\n" + table, flush=True)
print(f"Machine-readable report: {args.results_dir / (label + '.json')}", flush=True)
# Record all observations before returning a failing gate; never hide failures.
return (not report["input_errors"] and not report["inference_errors"]
and all(r["passed"] == r["count"] for r in report["divergence"].values())
and all(r["expectation_met"] for r in synth.values()))