diff --git a/parakeet_mlx/__init__.py b/parakeet_mlx/__init__.py index b1267dc..cbfa826 100644 --- a/parakeet_mlx/__init__.py +++ b/parakeet_mlx/__init__.py @@ -1,9 +1,9 @@ from parakeet_mlx.alignment import AlignedResult, AlignedSentence, AlignedToken from parakeet_mlx.parakeet import ( BaseParakeet, - DecodingConfig, ParakeetCTC, ParakeetCTCArgs, + ParakeetDecodingConfig, ParakeetRNNT, ParakeetRNNTArgs, ParakeetTDT, @@ -15,7 +15,7 @@ from parakeet_mlx.utils import from_pretrained __all__ = [ - "DecodingConfig", + "ParakeetDecodingConfig", "ParakeetTDTArgs", "ParakeetTDT", "ParakeetRNNT", diff --git a/parakeet_mlx/attention.py b/parakeet_mlx/attention.py index 7cfdc3c..a9879e7 100644 --- a/parakeet_mlx/attention.py +++ b/parakeet_mlx/attention.py @@ -44,7 +44,7 @@ def __call__( k, v = cache.update_and_fetch_kv(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) + o = o.transpose(0, 2, 1, 3).reshape(batch, q_seq, self.head_dim * self.n_head) return self.linear_out(o) @@ -534,6 +534,52 @@ def matmul_pv(self, prob: mx.array, v: mx.array, w: int) -> mx.array: return outputs[0] +# note that this has slight different scaling method than other encodings +class FixedPositionalEncoding(nn.Module): + def __init__( + self, + d_model: int, + max_len: int = 5000, + ): + 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) + self.calculate_pe() + + def calculate_pe(self): + positions = mx.arange(self.max_len, dtype=mx.float32) + positions = mx.expand_dims(positions, axis=1) + + div_term = mx.exp( + mx.arange(0, self.d_model, 2, dtype=mx.float32) + * -(math.log(10000.0) / self.d_model) + ) + pe = mx.zeros((self.max_len, 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) / self.scale + ) # we scale here! + + mx.eval(self._pe) + + def __call__(self, x: mx.array, offset: int = 0) -> mx.array: + input_len = x.shape[1] + + if offset + input_len > self.max_len: + self.max_len = offset + input_len + self.calculate_pe() + + pos_emb = self._pe[:, offset : offset + input_len, :].astype(x.dtype) + + return pos_emb + + class RelPositionalEncoding(nn.Module): def __init__( self, @@ -624,3 +670,24 @@ def __call__(self, x: mx.array, offset: int = 0) -> tuple[mx.array, mx.array]: pos_emb = self._pe[:, :end_idx].astype(x.dtype) return x, pos_emb + + +# utility +# thanks to mlx_lm +def create_causal_mask( + N: int, + offset: int = 0, + window_size: int | None = None, + lengths: mx.array | None = None, +): + rinds = mx.arange(offset + N) + linds = mx.arange(offset, offset + N) if offset else rinds + linds = linds[:, None] + rinds = rinds[None] + mask = linds >= rinds + if window_size is not None: + mask = mask & (linds <= rinds + window_size) + if lengths is not None: + lengths = lengths[:, None, None, None] + mask = mask & (rinds < lengths) + return mask diff --git a/parakeet_mlx/cache.py b/parakeet_mlx/cache.py index 0ad08eb..464b2f6 100644 --- a/parakeet_mlx/cache.py +++ b/parakeet_mlx/cache.py @@ -154,3 +154,49 @@ def update_and_fetch_conv(self, x: mx.array, padding: int = 0) -> mx.array: result = mx.pad(result, ((0, 0), (0, padding), (0, 0))) return result + + +class TransformerDecoderCache: + keys: mx.array | None + values: mx.array | None + + offset: int + step = 256 + + def __init__(self): + self.keys = None + self.values = None + self.conv = None + self.offset = 0 + + def update_and_fetch_kv( + self, keys: mx.array, values: mx.array + ) -> tuple[mx.array, mx.array]: + # k, v is [batch, head, seq, dim] + prev = self.offset + if ( + self.keys is None + or self.values is None + or (prev + keys.shape[2]) > self.keys.shape[2] + ): + B, H, S, D_KEYS = keys.shape + _, _, _, D_VALUES = values.shape + S_CACHE = ((self.step + S - 1) // self.step) * self.step + + new_k = mx.zeros((B, H, S_CACHE, D_KEYS), keys.dtype) + new_v = mx.zeros((B, H, S_CACHE, D_VALUES), keys.dtype) + + if self.keys is None or self.values is None: # type safety! + self.keys, self.values = new_k, new_v + else: + if prev % self.step != 0: + self.keys = self.keys[..., :prev, :] + self.values = self.values[..., :prev, :] + self.keys = mx.concatenate([self.keys, new_k], axis=2) + self.values = mx.concatenate([self.values, new_v], axis=2) + + self.offset += keys.shape[2] + self.keys[..., prev : self.offset, :] = keys + self.values[..., prev : self.offset, :] = values + + return self.keys[..., : self.offset, :], self.values[..., : self.offset, :] diff --git a/parakeet_mlx/canary.py b/parakeet_mlx/canary.py new file mode 100644 index 0000000..d6ffa7e --- /dev/null +++ b/parakeet_mlx/canary.py @@ -0,0 +1,419 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Optional, cast + +import mlx.core as mx +import mlx.nn as nn +from typing_extensions import Literal + +from parakeet_mlx.alignment import ( + AlignedResult, + AlignedToken, + merge_longest_common_subsequence, + merge_longest_contiguous, + sentences_to_result, + tokens_to_sentences, +) +from parakeet_mlx.audio import PreprocessArgs, get_logmel, load_audio +from parakeet_mlx.cache import TransformerDecoderCache +from parakeet_mlx.conformer import Conformer, ConformerArgs +from parakeet_mlx.tokenizer import CanaryTokenizer +from parakeet_mlx.transformer import ( + TransformerDecoder, + TransformerDecoderArgs, + TransformerHead, + TransformerHeadArgs, +) + + +@dataclass +class CanaryArgs: + preprocessor: PreprocessArgs + encoder: ConformerArgs + transf_decoder: TransformerDecoderArgs + head: TransformerHeadArgs + prompt_format: Literal["canary", "canary2"] + tokenizer: dict + + +@dataclass +class CanaryDecodingConfig: + decoding: Literal["greedy", "beam"] = "beam" + beam_size: int = 5 + temperature: float = 0.0 + max_length: int = 512 + + +class Canary(nn.Module): + """Canary model""" + + def __init__(self, args: CanaryArgs): + super().__init__() + + self.preprocessor_config = args.preprocessor + self.encoder_config = args.encoder + self.prompt_format: Literal["canary", "canary2"] = args.prompt_format + + self.tokenizer = CanaryTokenizer.from_data(args.tokenizer["data"]) + + self.encoder = Conformer(args.encoder) + self.transf_decoder = TransformerDecoder(args.transf_decoder) + self.head = TransformerHead(args.head) + + def transcribe( + self, + path: Path | str, + language: str, + timestamps: bool = False, + punctuation: bool = True, + *, + dtype: mx.Dtype = mx.bfloat16, + chunk_duration: Optional[float] = None, + overlap_duration: float = 15.0, + chunk_callback: Optional[Callable] = None, + ) -> AlignedResult | str: + audio_path = Path(path) + audio_data = load_audio(audio_path, self.preprocessor_config.sample_rate, dtype) + + if chunk_duration is None: + prompt_tokens = prompt( + self.tokenizer, + self.prompt_format, + language, + language, + punctuation, + timestamp=timestamps, + ) + mel = get_logmel(audio_data, self.preprocessor_config) + return self.generate(mel, [prompt_tokens])[0] + + audio_length_seconds = len(audio_data) / self.preprocessor_config.sample_rate + if audio_length_seconds <= chunk_duration: + prompt_tokens = prompt( + self.tokenizer, + self.prompt_format, + language, + language, + punctuation, + timestamp=timestamps, + ) + mel = get_logmel(audio_data, self.preprocessor_config) + return self.generate(mel, [prompt_tokens])[0] + + chunk_samples = int(chunk_duration * self.preprocessor_config.sample_rate) + overlap_samples = int(overlap_duration * self.preprocessor_config.sample_rate) + all_tokens = [] + previous_text = "" + + for start in range(0, len(audio_data), chunk_samples - overlap_samples): + end = min(start + chunk_samples, len(audio_data)) + + if chunk_callback is not None: + chunk_callback(end, len(audio_data)) + + if end - start < self.preprocessor_config.hop_length: + break + + if self.format == "canary2" and previous_text: + prompt_tokens = prompt( + self.tokenizer, + self.prompt_format, + language, + language, + punctuation, + context=previous_text, + timestamp=timestamps, + ) + else: + prompt_tokens = prompt( + self.tokenizer, + self.prompt_format, + language, + language, + punctuation, + timestamp=timestamps, + ) + + chunk_audio = audio_data[start:end] + chunk_mel = get_logmel(chunk_audio, self.preprocessor_config) + chunk_result = self.generate(chunk_mel, [prompt_tokens])[0] + + if chunk_result.text: + previous_text = chunk_result.text + + chunk_offset = start / self.preprocessor_config.sample_rate + for sentence in chunk_result.sentences: + for token in sentence.tokens: + token.start += chunk_offset + token.end = token.start + token.duration + + if all_tokens: + try: + all_tokens = merge_longest_contiguous( + all_tokens, + chunk_result.tokens, + overlap_duration=overlap_duration + if timestamps + else chunk_duration, + ) + except RuntimeError: + all_tokens = merge_longest_common_subsequence( + all_tokens, + chunk_result.tokens, + overlap_duration=overlap_duration + if timestamps + else chunk_duration, + ) + else: + all_tokens = chunk_result.tokens + + result = sentences_to_result(tokens_to_sentences(all_tokens)) + return result if timestamps else result.text + + def generate( + self, + mel: mx.array, + prompts: list[list[int]], + *, + decoding_config: CanaryDecodingConfig = CanaryDecodingConfig(), + ) -> list[AlignedResult]: + if len(mel.shape) == 2: + mel = mx.expand_dims(mel, 0) + features, lengths = self.encoder(mel) + mx.eval(features, lengths) + decoded = self.decode(features, prompts, lengths, config=decoding_config) + + def parse_time(token): + try: + s = self.tokenizer.decode([token]) + if s.startswith("<|") and s.endswith("|>"): + time_str = s[2:-2] + if time_str.isdigit(): + return ( + int(time_str) + * self.encoder_config.subsampling_factor + / self.preprocessor_config.sample_rate + * self.preprocessor_config.hop_length + ) + except (ValueError, AttributeError, KeyError): + pass + return None + + result = [] + for batch_idx, batch in enumerate(decoded): + aligned = [] + batch_length = float( + lengths[batch_idx] if batch_idx < len(lengths) else lengths[-1] + ) + max_time = ( + batch_length + * self.encoder_config.subsampling_factor + / self.preprocessor_config.sample_rate + * self.preprocessor_config.hop_length + ) + + timestamp_indices = {} + for i, t in enumerate(batch): + if t in self.tokenizer.special_tokens: + time = parse_time(t) + if time is not None: + timestamp_indices[i] = time + + for i, t in enumerate(batch): + if t not in self.tokenizer.special_tokens: + start = None + for j in range(i - 1, -1, -1): + if j in timestamp_indices: + start = timestamp_indices[j] + break + if start is None: + start = 0.0 + + end = None + for j in range(i + 1, len(batch)): + if j in timestamp_indices: + end = timestamp_indices[j] + break + if end is None: + end = ( + max_time + if i == len(batch) - 1 or not timestamp_indices + else start + ) + + aligned.append( + AlignedToken( + t, + self.tokenizer.decode([t]), + start, + max(0, end - start), + ) + ) + + result.append(sentences_to_result(tokens_to_sentences(aligned))) + + return result + + def decode( + self, + features: mx.array, + prompt: list[list[int]], + lengths: Optional[mx.array] = None, + *, + config: CanaryDecodingConfig = CanaryDecodingConfig(), + ) -> list[list[int]]: + if config.decoding == "greedy": + outputs = [] + for batch, p in enumerate(prompt): + tokens = [] + inputs = p.copy() + cache = [ + TransformerDecoderCache() + for _ in range(len(self.transf_decoder.layers)) + ] + + feat = features[batch : batch + 1] + if lengths is not None: + feat = feat[:, : int(lengths[batch])] + + while len(tokens) + len(p) < config.max_length: + logits = self.head( + self.transf_decoder(mx.array([inputs]), feat, cache=cache) + ) + next_token = cast(int, mx.argmax(logits[:, -1], axis=-1).item()) + + if next_token == self.tokenizer.eos_id: + break + + inputs = [next_token] + tokens.append(next_token) + + outputs.append(tokens) + return outputs + elif config.decoding == "beam": + outputs = [] + + for batch, p in enumerate(prompt): + # (tokens, inputs, score) + beams = [([], p.copy(), 0)] + cache = [ + TransformerDecoderCache() + for _ in range(len(self.transf_decoder.layers)) + ] + + feat = features[batch : batch + 1] + if lengths is not None: + feat = feat[:, : int(lengths[batch])] + + for _ in range(config.max_length - len(p)): + logits = self.head( + self.transf_decoder( + mx.array([beam[1] for beam in beams]), + mx.repeat(feat, len(beams), 0), + cache=cache, + ) + ) + logprobs = nn.log_softmax( + logits[:, -1] / max(config.temperature, 1e-8) + ) + accumulated_logprobs = logprobs.flatten() + mx.array( + [beam[2] for beam in beams for _ in range(logprobs.shape[1])] + ) + + indices = mx.argpartition(accumulated_logprobs, -config.beam_size)[ + -config.beam_size : + ] + beam_indices = indices // logprobs.shape[1] + token_indices = indices % logprobs.shape[1] + + # handle updates + for c in cache: + if c.keys is not None and c.values is not None: + c.keys = c.keys[beam_indices] + c.values = c.values[beam_indices] + beams = [ + ( + beams[int(beam_indices[i])][0] + [int(token_indices[i])], + [int(token_indices[i])], + float(accumulated_logprobs[indices[i]]), + ) + if beams[int(beam_indices[i])][1][0] != self.tokenizer.eos_id + else ( + beams[int(beam_indices[i])][0], + [self.tokenizer.eos_id], + beams[int(beam_indices[i])][2], + ) + for i in range(config.beam_size) + ] + + # exit condition + if all(beam[1][0] == self.tokenizer.eos_id for beam in beams): + beams = list(sorted(beams, key=lambda x: x[2], reverse=True)) + outputs.append(beams[0][0][:-1]) + break + + if len(outputs) < batch + 1: + # out of step + beams = list(sorted(beams, key=lambda x: x[2], reverse=True)) + eos_beams = list( + filter(lambda x: x[1][0] == self.tokenizer.eos_id, beams) + ) + if len(eos_beams) > 0: + outputs.append(eos_beams[0][0][:-1]) + else: + outputs.append(beams[0][0]) + + return outputs + + raise NotImplementedError + + +def prompt( + tokenizer: CanaryTokenizer, + prompt_format: Literal["canary", "canary2"], + source_lang: str, + target_lang: str, + punctuation: bool, + *, + context: str = "", + emotion: Literal["undefined", "neutral", "angry", "happy", "sad"] = "undefined", + inverse_normalization: bool = False, + timestamp: bool = False, + diarize: bool = False, +): + if prompt_format == "canary" and ( + len(context) > 0 + or emotion != "undefined" + or inverse_normalization is True + or timestamp is True + or diarize is True + ): + raise ValueError( + "`context`, `emotion`, `inverse_normalization`, `timestamp`, `diarize` are only supported in `canary2` prompt format." + ) + + src, tgt = f"<|{source_lang}|>", f"<|{target_lang}|>" + pnc = "<|pnc|>" if punctuation else "<|nopnc|>" + + if prompt_format == "canary": + task = "<|transcribe|>" if source_lang == target_lang else "<|translate|>" + prompt_text = f"<|startoftranscript|>{src}{task}{tgt}{pnc}" + return tokenizer.encode(prompt_text, lang_id="spl_tokens") + + emo = f"<|emo:{emotion}|>" + itn = "<|itn|>" if inverse_normalization else "<|noitn|>" + ts = "<|timestamp|>" if timestamp else "<|notimestamp|>" + dia = "<|diarize|>" if diarize else "<|nodiarize|>" + + if context: + ctx_tokens = tokenizer.encode(context, lang_id=target_lang) + prompt_tokens = tokenizer.encode( + f"<|startofcontext|><|startoftranscript|>{emo}{src}{tgt}{pnc}{itn}{ts}{dia}", + lang_id="spl_tokens", + ) + return prompt_tokens[:1] + ctx_tokens + prompt_tokens[1:] + + prompt_text = ( + f"<|startofcontext|><|startoftranscript|>{emo}{src}{tgt}{pnc}{itn}{ts}{dia}" + ) + return tokenizer.encode(prompt_text, lang_id="spl_tokens") diff --git a/parakeet_mlx/cli.py b/parakeet_mlx/cli.py index 9fe832b..d27cbd9 100644 --- a/parakeet_mlx/cli.py +++ b/parakeet_mlx/cli.py @@ -16,6 +16,7 @@ from typing_extensions import Annotated from parakeet_mlx import AlignedResult, AlignedSentence, AlignedToken, from_pretrained +from parakeet_mlx.parakeet import BaseParakeet app = typer.Typer(no_args_is_help=True) @@ -322,6 +323,9 @@ def transcribe( ) try: + if not isinstance(loaded_model, BaseParakeet): + return # TODO: HANDLE ME + result: AlignedResult = loaded_model.transcribe( audio_path, dtype=bfloat16 if not fp32 else float32, diff --git a/parakeet_mlx/conformer.py b/parakeet_mlx/conformer.py index bdc5475..8d21067 100644 --- a/parakeet_mlx/conformer.py +++ b/parakeet_mlx/conformer.py @@ -38,10 +38,16 @@ class ConformerArgs: class FeedForward(nn.Module): - def __init__(self, d_model: int, d_ff: int, use_bias: bool = True): + def __init__( + self, + d_model: int, + d_ff: int, + use_bias: bool = True, + activation: Literal["relu", "silu"] = "silu", + ): super().__init__() self.linear1 = nn.Linear(d_model, d_ff, bias=use_bias) - self.activation = nn.SiLU() + self.activation = nn.SiLU() if activation == "silu" else nn.ReLU() self.linear2 = nn.Linear(d_ff, d_model, bias=use_bias) def __call__(self, x: mx.array) -> mx.array: diff --git a/parakeet_mlx/parakeet.py b/parakeet_mlx/parakeet.py index b1288aa..8018aee 100644 --- a/parakeet_mlx/parakeet.py +++ b/parakeet_mlx/parakeet.py @@ -71,7 +71,7 @@ class ParakeetTDTCTCArgs(ParakeetTDTArgs): # API @dataclass -class DecodingConfig: +class ParakeetDecodingConfig: decoding: str = "greedy" @@ -88,7 +88,10 @@ def __init__(self, preprocess_args: PreprocessArgs, encoder_args: ConformerArgs) self.encoder = Conformer(encoder_args) def generate( - self, mel: mx.array, *, decoding_config: DecodingConfig = DecodingConfig() + self, + mel: mx.array, + *, + decoding_config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> list[AlignedResult]: """ Generate transcription results from the Parakeet model, handling batches and single input. @@ -96,9 +99,9 @@ def generate( mel (mx.array): Mel-spectrogram input with shape [batch, sequence, mel_dim] for batch processing or [sequence, mel_dim] for single input. - decoding_config (DecodingConfig, optional): + decoding_config (ParakeetDecodingConfig, optional): Configuration object that controls decoding behavior and - parameters for the generation process. Defaults to DecodingConfig(). + parameters for the generation process. Defaults to ParakeetDecodingConfig(). Returns: list[AlignedResult]: List of transcription results with aligned tokens and sentences, one for each input in the batch. @@ -198,7 +201,7 @@ def transcribe_stream( depth=1, *, keep_original_attention: bool = False, - decoding_config: DecodingConfig = DecodingConfig(), + decoding_config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> "StreamingParakeet": """ Create a StreamingParakeet object for real-time (streaming) inference. @@ -221,9 +224,9 @@ def transcribe_stream( keep_original_attention (bool, optional): Whether to preserve the original attention class during streaming inference. Defaults to False. (Will switch to local attention.) - decoding_config (DecodingConfig, optional): + decoding_config (ParakeetDecodingConfig, optional): Configuration object that controls decoding behavior - Defaults to DecodingConfig(). + Defaults to ParakeetDecodingConfig(). Returns: StreamingParakeet: A context manager for streaming inference. """ @@ -263,7 +266,7 @@ def decode( last_token: Optional[list[Optional[int]]] = None, hidden_state: Optional[list[Optional[tuple[mx.array, mx.array]]]] = None, *, - config: DecodingConfig = DecodingConfig(), + config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> tuple[list[list[AlignedToken]], list[Optional[tuple[mx.array, mx.array]]]]: """Run TDT decoder with features, optional length and decoder state. Outputs list[list[AlignedToken]] and updated hidden state""" assert config.decoding == "greedy", ( @@ -352,7 +355,10 @@ def decode( return results, hidden_state def generate( - self, mel: mx.array, *, decoding_config: DecodingConfig = DecodingConfig() + self, + mel: mx.array, + *, + decoding_config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> list[AlignedResult]: if len(mel.shape) == 2: mel = mx.expand_dims(mel, 0) @@ -391,7 +397,7 @@ def decode( last_token: Optional[list[Optional[int]]] = None, hidden_state: Optional[list[Optional[tuple[mx.array, mx.array]]]] = None, *, - config: DecodingConfig = DecodingConfig(), + config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> tuple[list[list[AlignedToken]], list[Optional[tuple[mx.array, mx.array]]]]: """Run TDT decoder with features, optional length and decoder state. Outputs list[list[AlignedToken]] and updated hidden state""" assert config.decoding == "greedy", ( @@ -472,7 +478,10 @@ def decode( return results, hidden_state def generate( - self, mel: mx.array, *, decoding_config: DecodingConfig = DecodingConfig() + self, + mel: mx.array, + *, + decoding_config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> list[AlignedResult]: if len(mel.shape) == 2: mel = mx.expand_dims(mel, 0) @@ -503,7 +512,7 @@ def decode( features: mx.array, lengths: mx.array, *, - config: DecodingConfig = DecodingConfig(), + config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> list[list[AlignedToken]]: """Run CTC decoder with features and lengths. Outputs list[list[AlignedToken]].""" B, S, *_ = features.shape @@ -596,7 +605,10 @@ def decode( return results def generate( - self, mel: mx.array, *, decoding_config: DecodingConfig = DecodingConfig() + self, + mel: mx.array, + *, + decoding_config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> list[AlignedResult]: if len(mel.shape) == 2: mel = mx.expand_dims(mel, 0) @@ -637,7 +649,7 @@ class StreamingParakeet: context_size: tuple[int, int] depth: int - decoding_config: DecodingConfig + decoding_config: ParakeetDecodingConfig keep_original_attention: bool = False def __init__( @@ -647,7 +659,7 @@ def __init__( depth: int = 1, *, keep_original_attention: bool = False, - decoding_config: DecodingConfig = DecodingConfig(), + decoding_config: ParakeetDecodingConfig = ParakeetDecodingConfig(), ) -> None: self.context_size = context_size self.depth = depth diff --git a/parakeet_mlx/tokenizer.py b/parakeet_mlx/tokenizer.py index 3758073..4e71c05 100644 --- a/parakeet_mlx/tokenizer.py +++ b/parakeet_mlx/tokenizer.py @@ -1,3 +1,85 @@ -# decode some tokens (might edit it if to support other varients) +import json +import re +from functools import cached_property + +from tokenizers import Tokenizer + + +# For parakeet: decode some tokens for parakeet def decode(tokens: list[int], vocabulary: list[str]): return "".join([vocabulary[token].replace("▁", " ") for token in tokens]) + + +# intersting approach.. +class CanaryTokenizer: + CANARY_BOS = "<|startoftranscript|>" + CANARY_EOS = "<|endoftext|>" + CANARY_PAD = "" + CANARY_NOSPEECH = "<|nospeech|>" + CANARY_PNC = "<|pnc|>" + CANARY_NOPNC = "<|nopnc|>" + CANARY2_BOCTX = "<|startofcontext|>" + + def __init__(self, tokenizers: dict[str, Tokenizer]): + self.tokenizers = tokenizers # dict with py 3.7+ is basically ordered dict + self.offsets = { + lang: sum(len(t.get_vocab()) for t in list(tokenizers.values())[:i]) + for i, lang in enumerate(tokenizers.keys()) + } + + mappings = [ + (local_id + self.offsets[lang], tokenizer, local_id) + for lang, tokenizer in tokenizers.items() + for local_id in tokenizer.get_vocab().values() + ] + self.lookup_tokenizer = {gid: tok for gid, tok, _ in mappings} + self.lookup_local = {gid: lid for gid, _, lid in mappings} + + self.special_tokens = { + token: local_id + self.offsets["spl_tokens"] + for token, local_id in tokenizers["spl_tokens"].get_vocab().items() + } + + @staticmethod + def from_data(tokenizer: dict[str, dict | str]): + return CanaryTokenizer( + { + i: Tokenizer.from_str(json.dumps(v) if isinstance(v, dict) else v) + for i, v in tokenizer.items() + } + ) + + def encode(self, text: str, lang_id: str) -> list[int]: + if lang_id == "spl_tokens": + return [ + self.special_tokens[token] for token in re.findall(r"<\|[^|]+\|>", text) + ] + + return [ + tid + self.offsets[lang_id] + for tid in self.tokenizers[lang_id].encode(text).ids + ] + + def decode(self, token_ids: list[int]) -> str: + pieces = [ + self.lookup_tokenizer[tid].decode([self.lookup_local[tid]]) + for tid in token_ids + if tid in self.lookup_tokenizer + ] + return "".join(pieces).replace("▁", " ") + + @cached_property + def eos_id(self) -> int: + return self.special_tokens[self.CANARY_EOS] + + @cached_property + def bos_id(self) -> int: + return self.special_tokens[self.CANARY_BOS] + + @cached_property + def nospeech_id(self) -> int: + return self.special_tokens[self.CANARY_NOSPEECH] + + @cached_property + def pad_id(self) -> int: + return self.special_tokens[self.CANARY_PAD] diff --git a/parakeet_mlx/transformer.py b/parakeet_mlx/transformer.py new file mode 100644 index 0000000..f7ec082 --- /dev/null +++ b/parakeet_mlx/transformer.py @@ -0,0 +1,152 @@ +from dataclasses import dataclass +from typing import Literal + +import mlx.core as mx +import mlx.nn as nn + +from parakeet_mlx.attention import ( + FixedPositionalEncoding, + MultiHeadAttention, + create_causal_mask, +) +from parakeet_mlx.cache import TransformerDecoderCache +from parakeet_mlx.conformer import FeedForward + + +@dataclass +class TransformerDecoderArgs: + vocab_size: int # num_classes in head + hidden_size: int + inner_size: int + num_layers: int + num_attention_heads: int + pre_ln: bool + hidden_act: Literal["relu"] + pre_ln_final_layer_norm: bool + learn_positional_encodings: bool + max_sequence_length: int + + +@dataclass +class TransformerHeadArgs: + num_layers: int + hidden_size: int + num_classes: int # this! + + +class TransformerDecoderBlock(nn.Module): + def __init__(self, args: TransformerDecoderArgs): + super().__init__() + + if args.pre_ln is False: + raise NotImplementedError( + "`pre_ln` = False for TransformerDecoder has not been implemented yet. Please open the issue in https://github.com/senstella/parakeet-mlx if you see this error." + ) + + self.layer_norm_1 = nn.LayerNorm(args.hidden_size, eps=1e-5) + self.first_sub_layer = MultiHeadAttention( + args.num_attention_heads, + args.hidden_size, + ) + self.layer_norm_2 = nn.LayerNorm(args.hidden_size, eps=1e-5) + self.second_sub_layer = MultiHeadAttention( + args.num_attention_heads, + args.hidden_size, + ) + self.layer_norm_3 = nn.LayerNorm(args.hidden_size, eps=1e-5) + self.third_sub_layer = FeedForward( + args.hidden_size, args.inner_size, activation=args.hidden_act + ) + + def __call__( + self, + x: mx.array, + xa: mx.array, + mask_x: mx.array | None = None, + mask_xa: mx.array | None = None, + cache: TransformerDecoderCache | None = None, + ) -> mx.array: + x_norm = self.layer_norm_1(x) + x = x + self.first_sub_layer(x_norm, x_norm, x_norm, mask=mask_x, cache=cache) + + x_norm = self.layer_norm_2(x) + x = x + self.second_sub_layer(x_norm, xa, xa, mask=mask_xa) + + x_norm = self.layer_norm_3(x) + x = x + self.third_sub_layer(x_norm) + + return x + + +class TransformerDecoder(nn.Module): + def __init__(self, args: TransformerDecoderArgs): + super().__init__() + + self.token_embedding = nn.Embedding( + args.vocab_size, args.hidden_size + ) # vocab_size is num_classes in head, kind of confusing naming + self.position_embedding = ( + nn.Embedding(args.max_sequence_length, args.hidden_size) + if args.learn_positional_encodings + else FixedPositionalEncoding( + args.hidden_size, max_len=args.max_sequence_length + ) + ) + self.embedding_layer_norm = nn.LayerNorm(args.hidden_size, eps=1e-5) + + self.layers = [TransformerDecoderBlock(args) for _ in range(args.num_layers)] + self.final_layer_norm = ( + nn.LayerNorm(args.hidden_size, eps=1e-5) + if args.pre_ln and args.pre_ln_final_layer_norm + else None + ) + + def __call__( + self, + x: mx.array, + xa: mx.array, + mask_x: mx.array | None = None, + mask_xa: mx.array | None = None, + cache: list[TransformerDecoderCache] | None = None, + ) -> mx.array: + # embedding + offset = 0 if cache is None else cache[0].offset + + x = self.token_embedding(x) + x = x + ( + self.position_embedding(x, offset=offset) + if isinstance(self.position_embedding, FixedPositionalEncoding) + else self.position_embedding(mx.arange(offset, offset + x.shape[1])) + ) + x = self.embedding_layer_norm(x) + + mask_x = ( + mask_x & create_causal_mask(x.shape[1], offset) + if mask_x is not None + else create_causal_mask(x.shape[1], offset) + ) + for i, layer in enumerate(self.layers): + x = layer( + x, xa, mask_x, mask_xa, cache=cache[i] if cache is not None else None + ) + + if self.final_layer_norm is not None: + x = self.final_layer_norm(x) + + return x + + +class TransformerHead(nn.Module): + def __init__(self, args: TransformerHeadArgs): + super().__init__() + + if args.num_layers != 1: + raise NotImplementedError( + "Classification head has non-supported layers. Please open an issue in https://github.com/senstella/parakeet-mlx" + ) + + self.classifier = nn.Linear(args.hidden_size, args.num_classes) + + def __call__(self, x: mx.array) -> mx.array: + x = self.classifier(x) + return x diff --git a/parakeet_mlx/utils.py b/parakeet_mlx/utils.py index 2d8663c..d380a53 100644 --- a/parakeet_mlx/utils.py +++ b/parakeet_mlx/utils.py @@ -6,6 +6,7 @@ from huggingface_hub import hf_hub_download from mlx.utils import tree_flatten, tree_unflatten +from parakeet_mlx.canary import Canary, CanaryArgs from parakeet_mlx.parakeet import ( BaseParakeet, ParakeetCTC, @@ -19,7 +20,7 @@ ) -def from_config(config: dict) -> BaseParakeet: +def from_config(config: dict) -> BaseParakeet | Canary: """Loads model from config (randomized weight)""" if ( config.get("target") @@ -48,6 +49,12 @@ def from_config(config: dict) -> BaseParakeet: ): cfg = from_dict(ParakeetCTCArgs, config) model = ParakeetCTC(cfg) + elif ( + config.get("target") + == "nemo.collections.asr.models.aed_multitask_models.EncDecMultiTaskModel" + ): + cfg = from_dict(CanaryArgs, config) + model = Canary(cfg) else: raise ValueError("Model is not supported yet!") @@ -58,7 +65,7 @@ def from_config(config: dict) -> BaseParakeet: def from_pretrained( hf_id_or_path: str, *, dtype: mx.Dtype = mx.bfloat16 -) -> BaseParakeet: +) -> BaseParakeet | Canary: """Loads model from Hugging Face or local directory""" try: config = json.load(open(hf_hub_download(hf_id_or_path, "config.json"), "r")) diff --git a/pyproject.toml b/pyproject.toml index b5d093f..58c9cdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "librosa>=0.11.0", "mlx>=0.22.1", "numpy>=2.2.5", + "tokenizers>=0.20.4", "typer>=0.15.3", ] license = "Apache-2.0"