Kuroko: weights, the file that runs them, settings, dev-split predictions, training log and card
879acfa | """Kuroko: of two speakers on separate channels, does the sound a speaker is making claim the floor? | |
| Kuroko hears both speakers as log-mel frames and answers 25 times a second for each of them (`ANSWERS`). An | |
| answer rests on the audio up to its own moment and on nothing later, but for the 0.6 ms a resampler reaches ahead | |
| at 48 kHz, so the answers can be used while a conversation unfolds. The write-up says how the answers become | |
| end-of-turn and interruption events: https://spero.patientdesk.ai/blog/kuroko | |
| import soundfile as sf | |
| import kuroko | |
| model = kuroko.load(".") # the directory that holds model.safetensors and config.json | |
| audio, rate = sf.read("conversation.wav", dtype="int16") # (samples, 2): one speaker on each channel | |
| answers = kuroko.answers(model, kuroko.features(audio, rate)) # (2, 11, n): speaker, answer, 25 a second | |
| claims = answers[:, kuroko.ANSWERS.index("claim")] | |
| This file is the model and what it needs to run, and nothing of how it was trained. It needs numpy, scipy, torch | |
| and safetensors. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from numpy.typing import NDArray | |
| from scipy.signal import resample_poly | |
| from torch import nn | |
| RATE = 16_000 # the sample rate the model hears | |
| HOP = 320 # 20 ms: one frame | |
| WINDOW = 2 * HOP # a frame holds the 40 ms that end with it | |
| FRAME_HZ = RATE // HOP | |
| ANSWER_HZ = FRAME_HZ // 2 # the model answers every other frame | |
| MELS = 80 | |
| LOW_HZ, HIGH_HZ = 60.0, 7_600.0 | |
| WINDOW_S = 40.0 # the most the model looks back | |
| AHEAD_S = ((0.0, 0.2), (0.2, 0.6), (0.6, 1.2), (1.2, 2.0)) | |
| ANSWERS = ( | |
| "voice", # the speaker is making a sound | |
| "claim", # that sound claims the floor: the answer the events are read from | |
| "hold", # the speaker holds the floor, pauses included | |
| *(f"{who} speaks {start:g} to {end:g} s ahead" for who in ("own", "partner") for start, end in AHEAD_S), | |
| ) | |
| _BLOCK = 20_000 # frames turned into spectra at once | |
| def mel_filters() -> NDArray[np.float32]: | |
| """Triangular filters (MELS, WINDOW // 2 + 1), evenly spaced on the mel scale from LOW_HZ to HIGH_HZ.""" | |
| low, high = (2595.0 * math.log10(1.0 + hz / 700.0) for hz in (LOW_HZ, HIGH_HZ)) | |
| edges = 700.0 * (10.0 ** (np.linspace(low, high, MELS + 2) / 2595.0) - 1.0) | |
| bins = np.fft.rfftfreq(WINDOW, 1.0 / RATE) | |
| rising = (bins[None] - edges[:-2, None]) / (edges[1:-1] - edges[:-2])[:, None] | |
| falling = (edges[2:, None] - bins[None]) / (edges[2:] - edges[1:-1])[:, None] | |
| filters: NDArray[np.float32] = np.clip(np.minimum(rising, falling), 0.0, None).astype(np.float32) | |
| return filters | |
| def log_mel(audio: NDArray[np.float32]) -> NDArray[np.float32]: | |
| """Log-mel frames (frames, MELS) of 16 kHz audio: frame k holds the 40 ms that end at sample (k + 1) * HOP.""" | |
| frames = len(audio) // HOP | |
| padded = np.concatenate([np.zeros(HOP, np.float32), audio[: frames * HOP]]) | |
| windows = np.lib.stride_tricks.sliding_window_view(padded, WINDOW)[::HOP] | |
| taper = np.hanning(WINDOW).astype(np.float32) | |
| filters = mel_filters() | |
| out = np.empty((frames, MELS), np.float32) | |
| for start in range(0, frames, _BLOCK): | |
| power = np.abs(np.fft.rfft(windows[start : start + _BLOCK] * taper, axis=1)) ** 2 | |
| out[start : start + _BLOCK] = np.log(power @ filters.T + 1e-10) | |
| return out | |
| def features(audio: NDArray[np.int16] | NDArray[np.floating], rate: int) -> NDArray[np.float16]: | |
| """The frames of a conversation given as (samples, 2), one speaker on each channel, at `rate` Hz, as 16-bit | |
| samples or as floats between -1 and 1: (2, frames, MELS), in the half precision the model learned from. | |
| The model was trained on 16-bit recordings at 48 kHz, brought to 16 kHz this way. Read as floats, audio of | |
| more than 16 bits gives answers that differ in the second decimal.""" | |
| scale = 32768.0 if audio.dtype == np.int16 else 1.0 | |
| step = math.gcd(RATE, rate) | |
| sides = [] | |
| for channel in (0, 1): | |
| heard = resample_poly(audio[:, channel].astype(np.float32) / scale, RATE // step, rate // step) | |
| sides.append(log_mel(heard.astype(np.float32)).astype(np.float16)) | |
| return np.stack(sides) | |
| def _turned(x: torch.Tensor, angle: torch.Tensor) -> torch.Tensor: | |
| """Rotary positions: each pair of `x`'s last dimension turned by its position's angle.""" | |
| first, second = x[..., 0::2], x[..., 1::2] | |
| cos, sin = angle.cos(), angle.sin() | |
| return torch.stack([first * cos - second * sin, first * sin + second * cos], dim=-1).flatten(-2) | |
| class Block(nn.Module): | |
| """One transformer layer whose attention sees the frames up to its own.""" | |
| def __init__(self, width: int, heads: int) -> None: | |
| super().__init__() | |
| self.heads = heads | |
| self.before_attention = nn.LayerNorm(width) | |
| self.query_key_value = nn.Linear(width, 3 * width) | |
| self.attended = nn.Linear(width, width) | |
| self.before_mlp = nn.LayerNorm(width) | |
| self.mlp = nn.Sequential(nn.Linear(width, 4 * width), nn.GELU(), nn.Linear(4 * width, width)) | |
| def forward(self, x: torch.Tensor, angle: torch.Tensor) -> torch.Tensor: | |
| batch, frames, width = x.shape | |
| parts = self.query_key_value(self.before_attention(x)).view(batch, frames, 3, self.heads, -1).transpose(1, 3) | |
| query, key, value = _turned(parts[:, :, 0], angle), _turned(parts[:, :, 1], angle), parts[:, :, 2] | |
| mixed = nn.functional.scaled_dot_product_attention(query, key, value, is_causal=True) | |
| x = x + self.attended(mixed.transpose(1, 2).reshape(batch, frames, width)) | |
| out: torch.Tensor = x + self.mlp(self.before_mlp(x)) | |
| return out | |
| class Kuroko(nn.Module): | |
| """Frames of both speakers in, a logit for each of `ANSWERS` and each speaker out, at half the frame rate. | |
| Both speakers go through the same weights, each once as "own" and once as "partner".""" | |
| mean: torch.Tensor | |
| std: torch.Tensor | |
| turns: torch.Tensor # how fast each pair of an attention head's dimensions turns with position | |
| def __init__(self, mels: int = MELS, width: int = 384, layers: int = 8, heads: int = 6) -> None: | |
| super().__init__() | |
| self.register_buffer("mean", torch.zeros(mels)) | |
| self.register_buffer("std", torch.ones(mels)) | |
| self.inputs = nn.Conv1d(mels, width, 5) | |
| self.halve = nn.Conv1d(width, width, 2, stride=2) | |
| self.join = nn.Linear(2 * width, width) | |
| self.blocks = nn.ModuleList(Block(width, heads) for _ in range(layers)) | |
| self.after = nn.LayerNorm(width) | |
| self.answers = nn.Linear(width, len(ANSWERS)) | |
| pairs = width // heads // 2 | |
| self.register_buffer("turns", 1.0 / 10_000 ** (torch.arange(pairs) / pairs), persistent=False) | |
| def forward(self, mel: torch.Tensor) -> torch.Tensor: | |
| """`mel` (batch, 2, frames, mels), from where the model starts to listen, gives logits | |
| (batch, 2, frames // 2, len(ANSWERS)). An answer sees its own two input frames and what came before.""" | |
| batch, speakers, frames, mels = mel.shape | |
| frames = frames // 2 * 2 | |
| heard = ((mel[:, :, :frames] - self.mean) / self.std).reshape(batch * speakers, frames, mels).transpose(1, 2) | |
| heard = nn.functional.gelu(self.inputs(nn.functional.pad(heard, (4, 0)))) # zeros before the audio | |
| own = nn.functional.gelu(self.halve(heard)).transpose(1, 2).reshape(batch, speakers, frames // 2, -1) | |
| x = self.join(torch.cat([own, own.flip(1)], dim=-1)).reshape(batch * speakers, frames // 2, -1) | |
| angle = torch.arange(frames // 2, device=x.device)[:, None] * self.turns | |
| for block in self.blocks: | |
| x = block(x, angle) | |
| logits: torch.Tensor = self.answers(self.after(x)).reshape(batch, speakers, frames // 2, len(ANSWERS)) | |
| return logits | |
| def load(directory: str | Path = ".", device: str | torch.device = "cpu") -> Kuroko: | |
| """The model whose `config.json` and `model.safetensors` are in a directory, ready to answer.""" | |
| from safetensors.torch import load_file | |
| config = json.loads((Path(directory) / "config.json").read_text()) | |
| model = Kuroko(config["mels"], config["width"], config["layers"], config["heads"]) | |
| model.load_state_dict(load_file(Path(directory) / "model.safetensors")) | |
| return model.to(device).eval() | |
| def answers(model: Kuroko, mel: NDArray[np.floating]) -> NDArray[np.float32]: | |
| """The model's answers for a whole conversation given as (2, frames, MELS): probabilities | |
| (2, len(ANSWERS), frames // 2), the answer at index k having heard the audio up to (k + 1) / ANSWER_HZ seconds. | |
| The model listens to stretches of `WINDOW_S` that start half a stretch apart, and of each but the first only | |
| the answers of its second half are kept: every answer then comes from a model that has listened for at least | |
| half a stretch, or from the conversation's start. | |
| Each stretch goes through the model by itself and in full precision. Batched or in half precision, an answer | |
| shifts in its third decimal with what else is in the batch, and with it the frame at which a threshold is | |
| crossed: the answers before a moment must not change with the audio after it.""" | |
| device = model.mean.device | |
| window = round(WINDOW_S * FRAME_HZ) | |
| half = window // 2 | |
| frames = mel.shape[1] // 2 * 2 | |
| out = np.zeros((2, len(ANSWERS), frames // 2), np.float32) | |
| for start in range(0, max(1, frames - half), half): | |
| stretch = np.zeros((1, 2, window, mel.shape[2]), np.float32) | |
| piece = mel[:, start : start + window] | |
| stretch[0, :, : piece.shape[1]] = piece | |
| scores = torch.sigmoid(model(torch.from_numpy(stretch).to(device))[0]).permute(0, 2, 1).cpu().numpy() | |
| keep = 0 if start == 0 else half // 2 # in answer frames | |
| end = min(frames, start + window) // 2 | |
| out[:, :, start // 2 + keep : end] = scores[:, :, keep : end - start // 2] | |
| return out | |