nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/mlx_audio/stt/models/parakeet/alignment.py

@@ -0,0 +1,248 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class AlignedToken:
+    id: int
+    text: str
+    start: float
+    duration: float
+    end: float = 0.0  # temporary
+
+    def __post_init__(self) -> None:
+        self.end = self.start + self.duration
+
+
+@dataclass
+class AlignedSentence:
+    text: str
+    tokens: list[AlignedToken]
+    start: float = 0.0  # temporary
+    end: float = 0.0  # temporary
+    duration: float = 0.0  # temporary
+
+    def __post_init__(self) -> None:
+        self.tokens = list(sorted(self.tokens, key=lambda x: x.start))
+        self.start = self.tokens[0].start
+        self.end = self.tokens[-1].end
+        self.duration = self.end - self.start
+
+
+@dataclass
+class AlignedResult:
+    text: str
+    sentences: list[AlignedSentence]
+
+    def __post_init__(self) -> None:
+        self.text = self.text.strip()
+
+
+def tokens_to_sentences(tokens: list[AlignedToken]) -> list[AlignedSentence]:
+    sentences = []
+    current_tokens = []
+
+    for idx, token in enumerate(tokens):
+        current_tokens.append(token)
+
+        # hacky, will fix
+        if (
+            "!" in token.text
+            or "?" in token.text
+            or "。" in token.text
+            or "？" in token.text
+            or "！" in token.text
+            or (
+                "." in token.text
+                and (idx == len(tokens) - 1 or " " in tokens[idx + 1].text)
+            )
+        ):  # type: ignore
+            sentence_text = "".join(t.text for t in current_tokens)
+            sentence = AlignedSentence(text=sentence_text, tokens=current_tokens)
+            sentences.append(sentence)
+
+            current_tokens = []
+
+    if current_tokens:
+        sentence_text = "".join(t.text for t in current_tokens)
+        sentence = AlignedSentence(text=sentence_text, tokens=current_tokens)
+        sentences.append(sentence)
+
+    return sentences
+
+
+def sentences_to_result(sentences: list[AlignedSentence]) -> AlignedResult:
+    return AlignedResult("".join(sentence.text for sentence in sentences), sentences)
+
+
+def merge_longest_contiguous(
+    a: list[AlignedToken],
+    b: list[AlignedToken],
+    *,
+    overlap_duration: float,
+):
+    if not a or not b:
+        return b if not a else a
+
+    a_end_time = a[-1].end
+    b_start_time = b[0].start
+
+    if a_end_time <= b_start_time:
+        return a + b
+
+    overlap_a = [token for token in a if token.end > b_start_time - overlap_duration]
+    overlap_b = [token for token in b if token.start < a_end_time + overlap_duration]
+
+    enough_pairs = len(overlap_a) // 2
+
+    if len(overlap_a) < 2 or len(overlap_b) < 2:
+        cutoff_time = (a_end_time + b_start_time) / 2
+        return [t for t in a if t.end <= cutoff_time] + [
+            t for t in b if t.start >= cutoff_time
+        ]
+
+    best_contiguous = []
+    for i in range(len(overlap_a)):
+        for j in range(len(overlap_b)):
+            if (
+                overlap_a[i].id == overlap_b[j].id
+                and abs(overlap_a[i].start - overlap_b[j].start) < overlap_duration / 2
+            ):
+                current = []
+                k, l_ = i, j
+                while (
+                    k < len(overlap_a)
+                    and l_ < len(overlap_b)
+                    and overlap_a[k].id == overlap_b[l_].id
+                    and abs(overlap_a[k].start - overlap_b[l_].start)
+                    < overlap_duration / 2
+                ):
+                    current.append((k, l_))
+                    k += 1
+                    l_ += 1
+
+                if len(current) > len(best_contiguous):
+                    best_contiguous = current
+
+    if len(best_contiguous) >= enough_pairs:
+        a_start_idx = len(a) - len(overlap_a)
+        lcs_indices_a = [a_start_idx + pair[0] for pair in best_contiguous]
+        lcs_indices_b = [pair[1] for pair in best_contiguous]
+
+        result = []
+        result.extend(a[: lcs_indices_a[0]])
+
+        for i in range(len(best_contiguous)):
+            idx_a = lcs_indices_a[i]
+            idx_b = lcs_indices_b[i]
+
+            result.append(a[idx_a])
+
+            if i < len(best_contiguous) - 1:
+                next_idx_a = lcs_indices_a[i + 1]
+                next_idx_b = lcs_indices_b[i + 1]
+
+                gap_tokens_a = a[idx_a + 1 : next_idx_a]
+                gap_tokens_b = b[idx_b + 1 : next_idx_b]
+
+                if len(gap_tokens_b) > len(gap_tokens_a):
+                    result.extend(gap_tokens_b)
+                else:
+                    result.extend(gap_tokens_a)
+
+        result.extend(b[lcs_indices_b[-1] + 1 :])
+        return result
+    else:
+        raise RuntimeError(f"No pairs exceeding {enough_pairs}")
+
+
+def merge_longest_common_subsequence(
+    a: list[AlignedToken],
+    b: list[AlignedToken],
+    *,
+    overlap_duration: float,
+):
+    if not a or not b:
+        return b if not a else a
+
+    a_end_time = a[-1].end
+    b_start_time = b[0].start
+
+    if a_end_time <= b_start_time:
+        return a + b
+
+    overlap_a = [token for token in a if token.end > b_start_time - overlap_duration]
+    overlap_b = [token for token in b if token.start < a_end_time + overlap_duration]
+
+    if len(overlap_a) < 2 or len(overlap_b) < 2:
+        cutoff_time = (a_end_time + b_start_time) / 2
+        return [t for t in a if t.end <= cutoff_time] + [
+            t for t in b if t.start >= cutoff_time
+        ]
+
+    dp = [[0 for _ in range(len(overlap_b) + 1)] for _ in range(len(overlap_a) + 1)]
+
+    for i in range(1, len(overlap_a) + 1):
+        for j in range(1, len(overlap_b) + 1):
+            if (
+                overlap_a[i - 1].id == overlap_b[j - 1].id
+                and abs(overlap_a[i - 1].start - overlap_b[j - 1].start)
+                < overlap_duration / 2
+            ):
+                dp[i][j] = dp[i - 1][j - 1] + 1
+            else:
+                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
+
+    lcs_pairs = []
+    i, j = len(overlap_a), len(overlap_b)
+
+    while i > 0 and j > 0:
+        if (
+            overlap_a[i - 1].id == overlap_b[j - 1].id
+            and abs(overlap_a[i - 1].start - overlap_b[j - 1].start)
+            < overlap_duration / 2
+        ):
+            lcs_pairs.append((i - 1, j - 1))
+            i -= 1
+            j -= 1
+        elif dp[i - 1][j] > dp[i][j - 1]:
+            i -= 1
+        else:
+            j -= 1
+
+    lcs_pairs.reverse()
+
+    if not lcs_pairs:
+        cutoff_time = (a_end_time + b_start_time) / 2
+        return [t for t in a if t.end <= cutoff_time] + [
+            t for t in b if t.start >= cutoff_time
+        ]
+
+    a_start_idx = len(a) - len(overlap_a)
+    lcs_indices_a = [a_start_idx + pair[0] for pair in lcs_pairs]
+    lcs_indices_b = [pair[1] for pair in lcs_pairs]
+
+    result = []
+
+    result.extend(a[: lcs_indices_a[0]])
+
+    for i in range(len(lcs_pairs)):
+        idx_a = lcs_indices_a[i]
+        idx_b = lcs_indices_b[i]
+
+        result.append(a[idx_a])
+
+        if i < len(lcs_pairs) - 1:
+            next_idx_a = lcs_indices_a[i + 1]
+            next_idx_b = lcs_indices_b[i + 1]
+
+            gap_tokens_a = a[idx_a + 1 : next_idx_a]
+            gap_tokens_b = b[idx_b + 1 : next_idx_b]
+
+            if len(gap_tokens_b) > len(gap_tokens_a):
+                result.extend(gap_tokens_b)
+            else:
+                result.extend(gap_tokens_a)
+
+    result.extend(b[lcs_indices_b[-1] + 1 :])
+
+    return result

nexaai/mlx_backend/mlx_audio/stt/models/parakeet/attention.py

@@ -0,0 +1,187 @@
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        n_head: int,
+        n_feat: int,
+        bias=True,
+    ):
+        super().__init__()
+
+        self.n_head = n_head
+        self.head_dim = n_feat // n_head
+        self.scale = self.head_dim**-0.5
+
+        self.linear_q = nn.Linear(n_feat, n_feat, bias=bias)
+        self.linear_k = nn.Linear(n_feat, n_feat, bias=bias)
+        self.linear_v = nn.Linear(n_feat, n_feat, bias=bias)
+        self.linear_out = nn.Linear(n_feat, n_feat, bias=bias)
+
+    def __call__(
+        self,
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        pos_emb: mx.array | None = None,
+        mask: mx.array | None = None,
+        cache=None,
+    ) -> mx.array:
+        q, k, v = self.linear_q(q), self.linear_k(k), self.linear_v(v)
+
+        batch, q_seq, _ = q.shape
+        _, k_seq, _ = k.shape
+
+        q = q.reshape(batch, q_seq, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
+        k = k.reshape(batch, k_seq, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
+        v = v.reshape(batch, k_seq, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
+
+        if cache:
+            k, v = cache.update_and_fetch(k, v)
+
+        o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
+        o = o.transpose(0, 2, 1, 3).reshape(batch, q_seq, self.n_feat)
+
+        return self.linear_out(o)
+
+
+class RelPositionMultiHeadAttention(MultiHeadAttention):
+    def __init__(
+        self,
+        n_head: int,
+        n_feat: int,
+        bias: bool = True,
+        pos_bias_u: mx.array | None = None,
+        pos_bias_v: mx.array | None = None,
+    ):
+        super().__init__(
+            n_head=n_head,
+            n_feat=n_feat,
+            bias=bias,
+        )
+
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+
+        if pos_bias_u is None:
+            self._pos_bias_u_init = mx.zeros((self.n_head, self.head_dim))
+        else:
+            self._pos_bias_u_init = pos_bias_u
+
+        if pos_bias_v is None:
+            self._pos_bias_v_init = mx.zeros((self.n_head, self.head_dim))
+        else:
+            self._pos_bias_v_init = pos_bias_v
+
+        self.pos_bias_u = self._pos_bias_u_init
+        self.pos_bias_v = self._pos_bias_v_init
+
+    def rel_shift(self, x: mx.array) -> mx.array:
+        B, H, Tq, pos_len = x.shape
+        padding = [(0, 0)] * (x.ndim - 1) + [(1, 0)]
+
+        x = mx.pad(x, padding)
+        x = x.reshape(B, H, pos_len + 1, Tq)
+        x = x[:, :, 1:, :]
+        x = x.reshape(B, H, Tq, pos_len)
+
+        return x
+
+    def __call__(
+        self,
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        pos_emb: mx.array | None = None,
+        mask: mx.array | None = None,
+        cache=None,
+    ) -> mx.array:
+        if pos_emb is None:
+            raise ValueError("pos_emb is necessary!")
+
+        q, k, v = self.linear_q(q), self.linear_k(k), self.linear_v(v)
+
+        p = self.linear_pos(pos_emb)  # p stands for position
+
+        batch, q_seq, _ = q.shape
+        _, k_seq, _ = k.shape
+        _, pos_len, _ = p.shape
+
+        q = q.reshape(batch, q_seq, self.n_head, self.head_dim)
+        q_u = (q + self.pos_bias_u).transpose(0, 2, 1, 3)
+        q_v = (q + self.pos_bias_v).transpose(0, 2, 1, 3)
+
+        k = k.reshape(batch, k_seq, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
+        v = v.reshape(batch, k_seq, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
+        p = p.reshape(batch, pos_len, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            k, v = cache.update_and_fetch(k, v)
+
+        matrix_bd = mx.matmul(q_v, p.swapaxes(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+        matrix_bd = matrix_bd[:, :, :, : k.shape[-2]] * self.scale
+
+        if mask is not None:
+            mask = mx.expand_dims(mask, 0)
+            matrix_bd[mask] = -mx.inf
+
+        o = mx.fast.scaled_dot_product_attention(
+            q_u, k, v, scale=self.scale, mask=matrix_bd
+        )
+        o = o.transpose(0, 2, 1, 3).reshape(batch, q_seq, -1)
+
+        return self.linear_out(o)
+
+
+class RelPositionalEncoding(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        max_len: int = 5000,
+        scale_input: bool = True,
+    ):
+        assert d_model % 2 == 0 and max_len > 0
+        super().__init__()
+
+        self.d_model = d_model
+        self.max_len = max_len
+        self.scale = math.sqrt(self.d_model) if scale_input else 1.0
+        self.calculate_pe()
+
+    def calculate_pe(self):
+        positions = mx.arange(self.max_len - 1, -self.max_len, -1, dtype=mx.int32)
+        positions = mx.expand_dims(positions, axis=1).astype(mx.float32)
+
+        div_term = mx.exp(
+            mx.arange(0, self.d_model, 2, dtype=mx.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe = mx.zeros((2 * self.max_len - 1, self.d_model), dtype=mx.float32)
+
+        pe[:, 0::2] = mx.sin(positions * div_term)
+        pe[:, 1::2] = mx.cos(positions * div_term)
+
+        self._pe = mx.expand_dims(pe, axis=0).astype(mx.float32)
+
+        mx.eval(self._pe)
+
+    def __call__(self, x: mx.array, offset: int = 0) -> tuple[mx.array, mx.array]:
+        input_len = x.shape[1] + offset
+
+        if input_len > self.max_len:
+            self.max_len = input_len + 1
+            self.calculate_pe()
+
+        x = x * self.scale
+
+        buffer_len = self._pe.shape[1]
+        start_idx = buffer_len // 2 - (input_len - 1)
+        end_idx = buffer_len // 2 + (input_len - 1) + 1
+
+        pos_emb = self._pe[:, start_idx:end_idx].astype(x.dtype)
+
+        return x, pos_emb

nexaai/mlx_backend/mlx_audio/stt/models/parakeet/audio.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+
+import mlx.core as mx
+
+from mlx_audio.utils import (
+    STR_TO_WINDOW_FN,
+    bartlett,
+    blackman,
+    hamming,
+    hanning,
+    mel_filters,
+    stft,
+)
+
+
+@dataclass
+class PreprocessArgs:
+    sample_rate: int
+    normalize: str
+    window_size: float
+    window_stride: float
+    window: str
+    features: int
+    n_fft: int
+    dither: float
+    pad_to: int = 0
+    pad_value: float = 0
+
+    @property
+    def win_length(self) -> int:
+        return int(self.window_size * self.sample_rate)
+
+    @property
+    def hop_length(self) -> int:
+        return int(self.window_stride * self.sample_rate)
+
+
+def log_mel_spectrogram(x: mx.array, args: PreprocessArgs) -> mx.array:
+    original_dtype = x.dtype
+
+    if args.pad_to > 0:
+        if x.shape[-1] < args.pad_to:
+            pad_length = args.pad_to - x.shape[-1]
+            x = mx.pad(x, ((0, pad_length),), constant_values=args.pad_value)
+
+    window_fn = STR_TO_WINDOW_FN.get(args.window, None)
+    window = window_fn(args.win_length) if window_fn else hanning(args.win_length)
+
+    # Ensure input is long enough for STFT (minimum n_fft samples required)
+    min_length = args.n_fft
+    if x.shape[-1] < min_length:
+        pad_length = min_length - x.shape[-1]
+        x = mx.pad(x, ((0, pad_length),), constant_values=args.pad_value)
+
+    x = stft(x, args.n_fft, args.hop_length, args.win_length, window)
+    x = mx.square(mx.abs(x)).astype(original_dtype)
+    filters = mel_filters(
+        args.sample_rate, args.n_fft, args.features, norm=args.normalize, mel_scale=None
+    )
+    x = filters.astype(x.dtype) @ x.T
+
+    x = mx.log(x + 1e-5)
+
+    if args.normalize == "per_feature":
+        mean = mx.mean(x, axis=1, keepdims=True)
+        std = mx.std(x, axis=1, keepdims=True)
+        normalized_mel = (x - mean) / (std + 1e-5)
+    else:
+        mean = mx.mean(x)
+        std = mx.std(x)
+        normalized_mel = (x - mean) / (std + 1e-5)
+
+    normalized_mel = normalized_mel.T
+    normalized_mel = mx.expand_dims(normalized_mel, axis=0)
+
+    return normalized_mel.astype(original_dtype)