| """Source parity: the authors' LaST-Net code (model.py LaSTNet + load_dataset.read_audio/crop_or_repeat + the infer.py |
| scoring loop, one clip at a time) vs this wrapper (batched), on J-SPAW_LA clips, at 4 s (Arena) and 6 s (authors' default). |
| |
| python source_parity.py <LaST-Net clone @ 3838f0c> [n_clips] |
| Deviations from running infer.py itself (both are loading-only, not math): |
| * eval.load_model does torch.load(best.pt, weights_only=False) -> not allowed here; the authors' LaSTNet instead |
| loads the same tensors from last_net_model.safetensors (= best.pt["model"], extract_state.py) with strict=True. |
| * fairseq.checkpoint_utils.load_model_ensemble_and_task([xlsr2_300m.pt]) fails on fairseq 0.12.2 (stale task cfg |
| keys) -> patched to build the same Wav2Vec2Model from the vendored xlsr2_300m model cfg; every tensor is then |
| overwritten by the strict load anyway. |
| TF32 off on both sides (FP32). |
| """ |
| import io |
| import os |
| import sys |
| import tempfile |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| import soundfile as sf |
| import torch |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, HERE) |
| src = sys.argv[1] |
| n = int(sys.argv[2]) if len(sys.argv) > 2 else 16 |
| import _net |
| import last_net as W |
| import fairseq |
|
|
| fairseq.checkpoint_utils.load_model_ensemble_and_task = lambda paths, **k: ([_net.build_xlsr()], None, None) |
| sys.path.append(src) |
| import model as src_model |
| from load_dataset import SAMPLE_RATE, crop_or_repeat, read_audio |
|
|
| torch.backends.cudnn.allow_tf32 = False |
| torch.backends.cuda.matmul.allow_tf32 = False |
| ref = src_model.LaSTNet(ssl_path="<patched>", temperature=0.8, uniform_floor=0.05, layer_dropout=0.10) |
| r = ref.load_state_dict(W.load_state(), strict=True) |
| print("source strict=True: missing", r.missing_keys, "unexpected", r.unexpected_keys) |
| ref = ref.cuda().eval() |
| ours = W.LaSTNetModel() |
| ours.load() |
|
|
| t = pq.read_table(os.path.join(HERE, "..", "J-SPAW_LA", "data", "test-00000-of-00001.parquet"), |
| columns=["audio", "label"]).to_pylist() |
| rows = [x for x in t if x["label"] == 0][: n // 2] + [x for x in t if x["label"] == 1][: n - n // 2] |
| auds = [sf.read(io.BytesIO(x["audio"]["bytes"]), dtype="float32")[0] for x in rows] |
| auds += [auds[0][:20000], auds[1][:5000], np.concatenate([auds[2]] * 3)] |
| tmp = tempfile.mkdtemp() |
| paths = [] |
| for k, a in enumerate(auds): |
| p = os.path.join(tmp, f"{k}.wav") |
| sf.write(p, a, 16000, subtype="FLOAT") |
| paths.append(p) |
| print("lens", [len(a) for a in auds]) |
| for seconds in (4, 6): |
| samples = seconds * SAMPLE_RATE |
| theirs = [] |
| with torch.inference_mode(): |
| for p in paths: |
| wav = read_audio(p) |
| available = min(len(wav), samples) / SAMPLE_RATE |
| batch = torch.from_numpy(crop_or_repeat(wav, samples)).unsqueeze(0).to("cuda") |
| duration = torch.tensor([available], device="cuda") |
| logits = ref(batch, duration).float() |
| theirs.append(logits.log_softmax(dim=1)[0, 1].item()) |
| theirs = np.array(theirs) |
| if seconds == 4: |
| mine = np.array(ours.score_batch(auds, [16000] * len(auds))) |
| else: |
| with torch.no_grad(): |
| xs, av = zip(*(_net.crop_np(a, samples) for a in auds)) |
| lg = ours.net(torch.from_numpy(np.stack(xs)).cuda(), torch.tensor(av, dtype=torch.float64).cuda()).float() |
| mine = lg.log_softmax(dim=1)[:, 1].cpu().numpy() |
| print(f"--- {seconds} s ({samples} samples), n={len(auds)}") |
| print("source", np.round(theirs, 4).tolist()) |
| print("ours ", np.round(np.asarray(mine, dtype=np.float64), 4).tolist()) |
| print(f"max_abs_diff_{seconds}s", float(np.max(np.abs(mine - theirs)))) |
|
|
