osobhandClaude Opus 4.7
479cd69bc3
- scripts/wavlm_sv_parity.py: Python-side reference embedder. Loads HF WavLMForXVector + Wav2Vec2FeatureExtractor and dumps a JSON fingerprint (cosine + per-utterance norm + first/last 8 elements) for comparison. - examples/wavlm_sv_demo gains --parity-json flag emitting the same fingerprint structure on the Rust side. Once the user has a Python env with transformers + torch installed, running both produces side-by-side JSON files for diffing — first-pass sanity check on whether our port matches HF numerically. We can't run the Python side from this Rust shell. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Phase 5d: WavLM-SV numerical parity check.
|
|
|
|
Compares the rtx-csm WavLM-SV port against the HF reference. Run once
|
|
the user has a Python env with `transformers` + `torch` + `safetensors`
|
|
available; emits a JSON report with the per-stage comparisons we can
|
|
diff against the Rust outputs.
|
|
|
|
Workflow:
|
|
pip install transformers torch torchaudio soundfile
|
|
|
|
# Generate the reference embeddings:
|
|
python3 scripts/wavlm_sv_parity.py \\
|
|
--wav-a /tmp/csm_24k.wav \\
|
|
--wav-b /tmp/test_24k.wav \\
|
|
--out /tmp/wavlm_sv_reference.json
|
|
|
|
# Then compare against the Rust outputs (TODO: extend wavlm_sv_demo
|
|
# to dump embeddings to JSON for diffing).
|
|
|
|
Why this lives in scripts/ rather than tests/: it requires PyTorch +
|
|
HF transformers as a heavy build-time dependency, and the comparison
|
|
itself is a one-shot numerical-validation step, not a regression gate.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--wav-a", required=True, type=Path)
|
|
parser.add_argument("--wav-b", required=True, type=Path)
|
|
parser.add_argument("--out", type=Path, default=Path("/tmp/wavlm_sv_reference.json"))
|
|
parser.add_argument("--model", default="microsoft/wavlm-base-plus-sv")
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
import torch
|
|
import soundfile as sf
|
|
from transformers import WavLMForXVector, AutoFeatureExtractor
|
|
except ImportError as e:
|
|
print(f"missing dependency: {e}", file=sys.stderr)
|
|
print(" pip install transformers torch soundfile", file=sys.stderr)
|
|
return 1
|
|
|
|
print(f"loading {args.model} (cached on first run)...")
|
|
extractor = AutoFeatureExtractor.from_pretrained(args.model)
|
|
model = WavLMForXVector.from_pretrained(args.model).eval()
|
|
|
|
def embed_path(path: Path) -> torch.Tensor:
|
|
# soundfile returns (n,) for mono WAV; resample if needed.
|
|
samples, sr = sf.read(str(path), dtype="float32")
|
|
if samples.ndim > 1:
|
|
samples = samples.mean(axis=1)
|
|
if sr != 16000:
|
|
import torchaudio
|
|
samples = torch.from_numpy(samples).unsqueeze(0)
|
|
samples = torchaudio.functional.resample(samples, sr, 16000)[0].numpy()
|
|
sr = 16000
|
|
inputs = extractor(samples, sampling_rate=sr, return_tensors="pt")
|
|
with torch.no_grad():
|
|
output = model(**inputs)
|
|
return output.embeddings[0] # (512,)
|
|
|
|
emb_a = embed_path(args.wav_a)
|
|
emb_b = embed_path(args.wav_b)
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
emb_a.unsqueeze(0), emb_b.unsqueeze(0)
|
|
).item()
|
|
|
|
report = {
|
|
"model": args.model,
|
|
"wav_a": str(args.wav_a),
|
|
"wav_b": str(args.wav_b),
|
|
"embedding_dim": emb_a.numel(),
|
|
"cosine_similarity_hf": cos,
|
|
"embedding_a_norm": emb_a.norm().item(),
|
|
"embedding_b_norm": emb_b.norm().item(),
|
|
# First and last 8 elements as a coarse fingerprint we can
|
|
# diff against the Rust output to catch axis/permute bugs.
|
|
"embedding_a_head": emb_a[:8].tolist(),
|
|
"embedding_a_tail": emb_a[-8:].tolist(),
|
|
"embedding_b_head": emb_b[:8].tolist(),
|
|
"embedding_b_tail": emb_b[-8:].tolist(),
|
|
}
|
|
args.out.write_text(json.dumps(report, indent=2))
|
|
print(f"wrote {args.out}")
|
|
print(f"HF cosine_similarity = {cos:.4f}")
|
|
print(f"|emb_a| = {report['embedding_a_norm']:.4f}, |emb_b| = {report['embedding_b_norm']:.4f}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|