| |
| |
| """Deterministic stand-ins for the two networks, for CPU tests of everything around them. |
| |
| Each output row is a pure function of that row's input (as measured for the real networks at the release's batch |
| shapes), so the release's `analyse` and the engine must build byte-identical records from them whatever the batching. |
| `FakeNN` is the engine-side backend; `patch_release` installs the same functions in the release's pipeline. |
| """ |
| import numpy as np |
|
|
| WIN, FRAMES, CLASSES, EMB = 256000, 799, 16, 256 |
| FRAMES_SW = [-0.00753125, 0.025, 0.02] |
| _F = np.arange(FRAMES, dtype=np.float64)[:, None] |
| _K = np.arange(CLASSES)[None, :] |
| _D = np.arange(EMB, dtype=np.float64) |
|
|
|
|
| def fake_segment(x): |
| x = np.asarray(x) |
| out = np.empty((x.shape[0], FRAMES, CLASSES), np.float32) |
| for r in range(x.shape[0]): |
| v = x[r, 0, :: WIN // 32].astype(np.float64) |
| a = float(np.abs(x[r, 0, ::997]).astype(np.float64).sum()) |
| logits = np.sin(_F * (0.011 + (a % 1.0) * 0.004) + _K * (1.3 + v[_K % 32] * 3.0)) * 5.0 |
| |
| logits[:, 0] += 6.0 * (a < 1e-3) |
| e = np.exp(logits - logits.max(1, keepdims=True)) |
| out[r] = (e / e.sum(1, keepdims=True)).astype(np.float32) |
| return out |
|
|
|
|
| def fake_embed(crops, weights): |
| crops, weights = np.asarray(crops), np.asarray(weights) |
| out = np.empty((crops.shape[0], EMB), np.float32) |
| for r in range(crops.shape[0]): |
| c = crops[r, 0, ::1000].astype(np.float64) |
| w = weights[r].astype(np.float64) |
| s = w.sum() |
| out[r] = (np.cos(c * 7.0 + _D * 0.1) + np.sin(_D * (1.0 + (s % 5.0))) + w[::4][:EMB].sum() * 1e-3).astype(np.float32) |
| return out |
|
|
|
|
| class FakeNN: |
| """Engine backend with the stand-in networks; records every call's shape.""" |
|
|
| def __init__(self, devices=(0, 0, 1)): |
| self._slots = list(devices) |
| self.calls = [] |
|
|
| def slot_devices(self): |
| return self._slots |
|
|
| def info(self): |
| return dict(frames_sw=FRAMES_SW) |
|
|
| def segment(self, slot, x): |
| self.calls.append(("segment", x.shape)) |
| return fake_segment(x) |
|
|
| def embed(self, slot, crops, weights): |
| self.calls.append(("embed", crops.shape, weights.shape)) |
| return fake_embed(crops, weights) |
|
|
| def embed_repacked(self, slot, windows, idx, weights): |
| self.calls.append(("embed_repacked", windows.shape, idx.shape, weights.shape)) |
| return np.concatenate([fake_embed(windows[idx[h]], weights[h]) for h in range(idx.shape[0])]) |
|
|
| def resample(self, audio, sr): |
| import torch |
| import torchaudio.functional as F |
| torch.set_num_threads(2) |
| return F.resample(torch.from_numpy(audio).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32) |
|
|
|
|
| def patch_release(model): |
| """Replace the networks inside a pipeline built by the release's `load_bundle`.""" |
| pipe = model["pipeline"] |
| pipe._segmentation.infer = lambda chunks, soft=False: fake_segment(chunks.numpy()) |
|
|
| class Emb: |
| sample_rate = 16000 |
| min_num_samples = 400 |
| dimension = EMB |
| metric = "cosine" |
|
|
| def __call__(self, waveforms, masks=None): |
| return fake_embed(waveforms.numpy(), masks.numpy()) |
|
|
| pipe._embedding = Emb() |
|
|
