| import mlx.core as mx |
| import mlx.nn as nn |
| from mlx_lm.models.base import BaseModelArgs |
| from mlx_lm.utils import load |
|
|
| from subsampling import ConvSubsampling |
| from modules import ConformerLayer |
| from attention import RelPositionalEncoding |
|
|
| class FastConformerEncoder(nn.Module): |
| def __init__( |
| self, |
| feat_in: int = 80, |
| n_layers: int = 18, |
| d_model: int = 512, |
| ff_expansion_factor: int = 4, |
| n_heads: int = 8, |
| conv_kernel_size: int = 9, |
| dropout: float = 0.1, |
| ): |
| super().__init__() |
| self.d_model = d_model |
| |
| self.pre_encode = ConvSubsampling( |
| subsampling='dw_striding', |
| subsampling_factor=8, |
| feat_in=feat_in, |
| feat_out=d_model, |
| conv_channels=256 |
| ) |
| |
| self.pos_enc = RelPositionalEncoding(d_model=d_model, max_len=5000) |
| |
| d_ff = d_model * ff_expansion_factor |
| self.layers = [ |
| ConformerLayer( |
| d_model=d_model, |
| d_ff=d_ff, |
| n_heads=n_heads, |
| conv_kernel_size=conv_kernel_size, |
| dropout=dropout |
| ) for _ in range(n_layers) |
| ] |
|
|
| def __call__(self, x, lengths=None): |
| |
| if lengths is None: |
| lengths = mx.array([x.shape[1]] * x.shape[0]) |
| |
| x, lengths = self.pre_encode(x, lengths) |
| |
| |
| x, pos_emb = self.pos_enc(x) |
| |
| for layer in self.layers: |
| x = layer(x, pos_emb=pos_emb) |
| |
| return x, lengths |
|
|
| class CanaryModel(nn.Module): |
| """ |
| Hybrid ASR-LLM Model connecting FastConformer to Qwen via a linear projection. |
| """ |
| def __init__(self, encoder: FastConformerEncoder, llm_dim: int): |
| super().__init__() |
| self.encoder = encoder |
| |
| |
| self.audio_encoder_proj = nn.Linear(1024, llm_dim) |
| |
| def encode_audio(self, audio_features, lengths=None): |
| """ |
| Passes audio features through the encoder and projects them to the LLM embedding space. |
| """ |
| encoded_audio, lengths = self.encoder(audio_features, lengths) |
| audio_embeds = self.audio_encoder_proj(encoded_audio) |
| return audio_embeds, lengths |
|
|
| |
| |
|
|
