| |
| |
| |
| import io, os, sys |
| import numpy as np, pyarrow.parquet as pq, soundfile as sf, torch |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, HERE) |
| W = __import__([m[:-3] for m in os.listdir(HERE) if m.startswith("wst_graph_q8") and m.endswith(".py")][0]) |
| cls = [getattr(W, n) for n in dir(W) if n.startswith("WSTGraphQ8")][0] |
| m = cls(); m.load() |
| ref = torch.jit.load(W._CKPT, map_location="cuda:0").eval() |
|
|
| d = os.path.join(HERE, "..", "J-SPAW_LA", "data") |
| t = pq.read_table(os.path.join(d, "test-00000-of-00001.parquet"), columns=["audio"]).slice(0, 16).column("audio").to_pylist() |
| auds = [sf.read(io.BytesIO(x["bytes"]), dtype="float32", always_2d=True)[0].mean(axis=1).astype(np.float32) for x in t] |
| auds.append(auds[0][:20000]); auds.append(np.concatenate([auds[1]] * 4)) |
| ours = np.array(m.score_batch(auds, [16000] * len(auds))) |
|
|
|
|
| def card(audio_list): |
| waveforms, lengths = [], [] |
| for audio in audio_list: |
| length = min(len(audio), 64000) |
| waveform = torch.zeros(64000, dtype=torch.float32) |
| waveform[:length] = torch.from_numpy(audio[:length]) |
| waveforms.append(waveform); lengths.append(length) |
| input_values = torch.stack(waveforms).to("cuda:0") |
| input_length = torch.tensor(lengths, dtype=torch.int64, device="cuda:0") |
| with torch.inference_mode(): |
| logits = ref(input_values, input_length) |
| return (logits[:, 0] - logits[:, 1]).cpu().numpy() |
|
|
|
|
| theirs_single = np.concatenate([card([a]) for a in auds]) |
| theirs_batch = card(auds) |
| print("lens", [len(a) for a in auds]) |
| print("ours ", np.round(ours, 4)) |
| print("card ", np.round(theirs_single, 4)) |
| print("max_abs_diff vs card per-clip", float(np.max(np.abs(ours - theirs_single)))) |
| print("max_abs_diff vs card same batch", float(np.max(np.abs(ours - theirs_batch)))) |
|
|
