Upload 9 files

added_tokens.json

@@ -0,0 +1,5 @@
+{
+    "</s>": 35,
+    "<s>": 34
+  }
+  

config.json

@@ -0,0 +1,73 @@
+{   
+    "auto_map":{
+        "AutoConfig": "gigaam_transformers.GigaAMConfig",
+        "AutoModel": "gigaam_transformers.GigaAMRNNTHF",
+        "AutoProcessor": "gigaam_transformers.GigaAMProcessor",
+        "AutoTokenizer": "gigaam_transformers.GigaAMTokenizer",
+        "AutoFeatureExtractor": "gigaam_transformers.GigaAMFeatureExtractor"
+    },
+
+    "encoder": {
+        "feat_in": 64,
+        "n_layers": 16,
+        "d_model": 768,
+        "subsampling_factor": 4,
+        "ff_expansion_factor": 4,
+        "self_attention_model": "rotary",
+        "pos_emb_max_len": 5000,
+        "n_heads": 16,
+        "conv_kernel_size": 31,
+        "flash_attn": false
+    },
+    "head": {
+        "decoder": {
+            "pred_hidden": 320,
+            "pred_rnn_layers": 1,
+            "num_classes": 34
+            },
+        "joint": {
+            "enc_hidden": 768,
+            "pred_hidden": 320,
+            "joint_hidden": 320,
+            "num_classes": 34
+            }
+    },
+    "labels": [
+        " ",
+        "а",
+        "б",
+        "в",
+        "г",
+        "д",
+        "е",
+        "ж",
+        "з",
+        "и",
+        "й",
+        "к",
+        "л",
+        "м",
+        "н",
+        "о",
+        "п",
+        "р",
+        "с",
+        "т",
+        "у",
+        "ф",
+        "х",
+        "ц",
+        "ч",
+        "ш",
+        "щ",
+        "ъ",
+        "ы",
+        "ь",
+        "э",
+        "ю",
+        "я"
+    ],
+    "blank_id": 33,
+    "max_symbols": 1000,
+    "model_type": "gigaam-rnnt"
+}

encoder.py

@@ -0,0 +1,604 @@
+"""Copied from https://github.com/salute-developers/GigaAM/blob/main/gigaam/encoder.py"""
+import math
+from abc import ABC, abstractmethod
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+
+# try:
+#     from flash_attn import flash_attn_func
+
+#     IMPORT_FLASH = True
+# except Exception as err:
+#     IMPORT_FLASH = False
+#     IMPORT_FLASH_ERR = err
+
+IMPORT_FLASH = False
+IMPORT_FLASH_ERR = "Flash Attention not installed."
+
+# from .utils import apply_masked_flash_attn, apply_rotary_pos_emb
+
+
+def rtt_half(x: Tensor) -> Tensor:
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat([-x2, x1], dim=x1.ndim - 1)
+
+
+def apply_rotary_pos_emb(
+    q: Tensor, k: Tensor, cos: Tensor, sin: Tensor, offset: int = 0
+) -> Tuple[Tensor, Tensor]:
+    """
+    Applies Rotary Position Embeddings to query and key tensors.
+    """
+    cos, sin = (
+        cos[offset : q.shape[0] + offset, ...],
+        sin[offset : q.shape[0] + offset, ...],
+    )
+    return (q * cos) + (rtt_half(q) * sin), (k * cos) + (rtt_half(k) * sin)
+
+
+# def apply_masked_flash_attn(
+#     q: Tensor,
+#     k: Tensor,
+#     v: Tensor,
+#     mask: Tensor,
+#     h: int,
+#     d_k: int,
+# ) -> Tensor:
+#     """
+#     Applies Flash Attention with padding masks.
+#     """
+
+#     from einops import rearrange
+#     from flash_attn import flash_attn_varlen_func
+#     from flash_attn.bert_padding import pad_input, unpad_input
+
+#     pad_mask = ~mask[:, 0, :]
+#     b, t = pad_mask.shape
+#     q = q.view(b, t, h * d_k)
+#     k = k.view(b, t, h * d_k)
+#     v = v.view(b, t, h * d_k)
+
+#     q_unpad, indices_q, _, max_seqlen_q = unpad_input(q, pad_mask)[:4]
+#     q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
+
+#     k_unpad = unpad_input(k, pad_mask)[0]
+#     k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
+
+#     v_unpad = unpad_input(v, pad_mask)[0]
+#     v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
+
+#     lengths_q = pad_mask.sum(1).to(torch.int32).to(q.device)
+#     cu_seqlens_q = F.pad(lengths_q.cumsum(0), (1, 0), value=0).to(torch.int32)
+#     max_seqlen_q = torch.max(lengths_q)
+
+#     output_unpad = flash_attn_varlen_func(
+#         q_unpad,
+#         k_unpad,
+#         v_unpad,
+#         cu_seqlens_q,
+#         cu_seqlens_q,
+#         max_seqlen_q,
+#         max_seqlen_q,
+#     )
+
+#     scores = pad_input(
+#         rearrange(output_unpad, "nnz h d -> nnz (h d)"),
+#         indices_q,
+#         b,
+#         t,
+#     )
+
+#     return scores
+
+
+class StridingSubsampling(nn.Module):
+    """
+    Strided Subsampling layer used to reduce the sequence length.
+    """
+
+    def __init__(
+        self,
+        subsampling_factor: int,
+        feat_in: int,
+        feat_out: int,
+        conv_channels: int,
+    ):
+        super().__init__()
+        self._sampling_num = int(math.log(subsampling_factor, 2))
+        self._stride = 2
+        self._kernel_size = 3
+        self._padding = (self._kernel_size - 1) // 2
+
+        layers: List[nn.Module] = []
+        in_channels = 1
+        for _ in range(self._sampling_num):
+            layers.append(
+                torch.nn.Conv2d(
+                    in_channels=in_channels,
+                    out_channels=conv_channels,
+                    kernel_size=self._kernel_size,
+                    stride=self._stride,
+                    padding=self._padding,
+                )
+            )
+            layers.append(nn.ReLU())
+            in_channels = conv_channels
+
+        out_length = self.calc_output_length(torch.tensor(feat_in))
+        self.out = torch.nn.Linear(conv_channels * int(out_length), feat_out)
+        self.conv = torch.nn.Sequential(*layers)
+
+    def calc_output_length(self, lengths: Tensor) -> Tensor:
+        """
+        Calculates the output length after applying the subsampling.
+        """
+        lengths = lengths.to(torch.float)
+        add_pad = 2 * self._padding - self._kernel_size
+        for _ in range(self._sampling_num):
+            lengths = torch.div(lengths + add_pad, self._stride) + 1.0
+            lengths = torch.floor(lengths)
+        return lengths.to(dtype=torch.int)
+
+    def forward(self, x: Tensor, lengths: Tensor) -> Tuple[Tensor, Tensor]:
+        x = self.conv(x.unsqueeze(1))
+        b, _, t, _ = x.size()
+        x = self.out(x.transpose(1, 2).reshape(b, t, -1))
+        return x, self.calc_output_length(lengths)
+
+
+class MultiHeadAttention(nn.Module, ABC):
+    """
+    Base class of Multi-Head Attention Mechanisms.
+    """
+
+    def __init__(self, n_head: int, n_feat: int, flash_attn=False):
+        super().__init__()
+        assert n_feat % n_head == 0
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.flash_attn = flash_attn
+        if self.flash_attn and not IMPORT_FLASH:
+            raise RuntimeError(
+                f"flash_attn_func was imported with err {IMPORT_FLASH_ERR}. "
+                "Please install flash_attn or use --no_flash flag. "
+                "If you have already done this, "
+                "--force-reinstall flag might be useful"
+            )
+
+    def forward_qkv(
+        self, query: Tensor, key: Tensor, value: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Projects the inputs into queries, keys, and values for multi-head attention.
+        """
+        b = query.size(0)
+        q = self.linear_q(query).view(b, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(b, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(b, -1, self.h, self.d_k)
+        if self.flash_attn:
+            return q, k, v
+        return q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+    def forward_attention(
+        self, value: Tensor, scores: Tensor, mask: Optional[Tensor]
+    ) -> Tensor:
+        """
+        Computes the scaled dot-product attention given the projected values and scores.
+        """
+        b = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+            scores = scores.masked_fill(mask, -10000.0)
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
+        else:
+            attn = torch.softmax(scores, dim=-1)
+        x = torch.matmul(attn, value)
+        x = x.transpose(1, 2).reshape(b, -1, self.h * self.d_k)
+        return self.linear_out(x)
+
+
+class RelPositionMultiHeadAttention(MultiHeadAttention):
+    """
+    Relative Position Multi-Head Attention module.
+    """
+
+    def __init__(self, n_head: int, n_feat: int):
+        super().__init__(n_head, n_feat)
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        self.pos_bias_u = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+
+    def rel_shift(self, x: Tensor) -> Tensor:
+        b, h, qlen, pos_len = x.size()
+        x = torch.nn.functional.pad(x, pad=(1, 0))
+        x = x.view(b, h, -1, qlen)
+        return x[:, :, 1:].view(b, h, qlen, pos_len)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        pos_emb: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)
+        p = self.linear_pos(pos_emb)
+        p = p.view(pos_emb.shape[0], -1, self.h, self.d_k).transpose(1, 2)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+        matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
+        scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask)
+
+
+class RotaryPositionMultiHeadAttention(MultiHeadAttention):
+    """
+    Rotary Position Multi-Head Attention module.
+    """
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        pos_emb: List[Tensor],
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        b, t, _ = value.size()
+        query = query.transpose(0, 1).view(t, b, self.h, self.d_k)
+        key = key.transpose(0, 1).view(t, b, self.h, self.d_k)
+        value = value.transpose(0, 1).view(t, b, self.h, self.d_k)
+
+        cos, sin = pos_emb
+        query, key = apply_rotary_pos_emb(query, key, cos, sin, offset=0)
+
+        q, k, v = self.forward_qkv(
+            query.view(t, b, self.h * self.d_k).transpose(0, 1),
+            key.view(t, b, self.h * self.d_k).transpose(0, 1),
+            value.view(t, b, self.h * self.d_k).transpose(0, 1),
+        )
+
+        # if not self.flash_attn:
+        scores = torch.matmul(q, k.transpose(-2, -1) / math.sqrt(self.d_k))
+        out = self.forward_attention(v, scores, mask)
+        # else:
+        #     if mask is None:
+        #         scores = flash_attn_func(q, k, v)
+        #     else:
+        #         scores = apply_masked_flash_attn(q, k, v, mask, self.h, self.d_k)
+
+        #     scores = scores.view(b, -1, self.h * self.d_k)
+        #     out = self.linear_out(scores)
+
+        return out
+
+
+class PositionalEncoding(nn.Module, ABC):
+    """
+    Base class of Positional Encodings.
+    """
+
+    def __init__(self, dim: int, base: int):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+
+    @abstractmethod
+    def create_pe(self, length: int, device: torch.device) -> Optional[Tensor]:
+        pass
+
+    def extend_pe(self, length: int, device: torch.device):
+        """
+        Extends the positional encoding buffer to process longer sequences.
+        """
+        pe = self.create_pe(length, device)
+        if pe is None:
+            return
+        if hasattr(self, "pe"):
+            self.pe = pe
+        else:
+            self.register_buffer("pe", pe, persistent=False)
+
+
+class RelPositionalEmbedding(PositionalEncoding):
+    """
+    Relative Positional Embedding module.
+    """
+
+    def create_pe(self, length: int, device: torch.device) -> Optional[Tensor]:
+        """
+        Creates the relative positional encoding matrix.
+        """
+        if hasattr(self, "pe") and self.pe.shape[1] >= 2 * length - 1:
+            return None
+        positions = torch.arange(length - 1, -length, -1, device=device).unsqueeze(1)
+        pos_length = positions.size(0)
+        pe = torch.zeros(pos_length, self.dim, device=positions.device)
+        div_term = torch.exp(
+            torch.arange(0, self.dim, 2, device=pe.device)
+            * -(math.log(10000.0) / self.dim)
+        )
+        pe[:, 0::2] = torch.sin(positions * div_term)
+        pe[:, 1::2] = torch.cos(positions * div_term)
+        return pe.unsqueeze(0)
+
+    def forward(self, x: torch.Tensor) -> Tuple[Tensor, Tensor]:
+        input_len = x.size(1)
+        center_pos = self.pe.size(1) // 2 + 1
+        start_pos = center_pos - input_len
+        end_pos = center_pos + input_len - 1
+        return x, self.pe[:, start_pos:end_pos]
+
+
+class RotaryPositionalEmbedding(PositionalEncoding):
+    """
+    Rotary Positional Embedding module.
+    """
+
+    def create_pe(self, length: int, device: torch.device) -> Optional[Tensor]:
+        """
+        Creates or extends the rotary positional encoding matrix.
+        """
+        if hasattr(self, "pe") and self.pe.size(0) >= 2 * length:
+            return None
+        positions = torch.arange(0, length, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)
+        )
+        t = torch.arange(length, device=positions.device).type_as(inv_freq)
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1).to(positions.device)
+        return torch.cat([emb.cos()[:, None, None, :], emb.sin()[:, None, None, :]])
+
+    def forward(self, x: torch.Tensor) -> Tuple[Tensor, List[Tensor]]:
+        cos_emb = self.pe[0 : x.shape[1]]
+        half_pe = self.pe.shape[0] // 2
+        sin_emb = self.pe[half_pe : half_pe + x.shape[1]]
+        return x, [cos_emb, sin_emb]
+
+
+class ConformerConvolution(nn.Module):
+    """
+    Conformer Convolution module.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        kernel_size: int,
+    ):
+        super().__init__()
+        assert (kernel_size - 1) % 2 == 0
+        self.pointwise_conv1 = nn.Conv1d(d_model, d_model * 2, kernel_size=1)
+        self.depthwise_conv = nn.Conv1d(
+            in_channels=d_model,
+            out_channels=d_model,
+            kernel_size=kernel_size,
+            padding=(kernel_size - 1) // 2,
+            groups=d_model,
+            bias=True,
+        )
+        self.batch_norm = nn.BatchNorm1d(d_model)
+        self.activation = nn.SiLU()
+        self.pointwise_conv2 = nn.Conv1d(d_model, d_model, kernel_size=1)
+
+    def forward(self, x: Tensor, pad_mask: Optional[Tensor] = None) -> Tensor:
+        x = x.transpose(1, 2)
+        x = self.pointwise_conv1(x)
+        x = nn.functional.glu(x, dim=1)
+        if pad_mask is not None:
+            x = x.masked_fill(pad_mask.unsqueeze(1), 0.0)
+        x = self.depthwise_conv(x)
+        x = self.batch_norm(x)
+        x = self.activation(x)
+        x = self.pointwise_conv2(x)
+        return x.transpose(1, 2)
+
+
+class ConformerFeedForward(nn.Module):
+    """
+    Conformer Feed Forward module.
+    """
+
+    def __init__(self, d_model: int, d_ff: int, use_bias=True):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff, bias=use_bias)
+        self.activation = nn.SiLU()
+        self.linear2 = nn.Linear(d_ff, d_model, bias=use_bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.linear2(self.activation(self.linear1(x)))
+
+
+class ConformerLayer(nn.Module):
+    """
+    Conformer Layer module.
+    This module combines several submodules including feed forward networks,
+    depthwise separable convolution, and multi-head self-attention
+    to form a single Conformer block.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int,
+        self_attention_model: str,
+        n_heads: int = 16,
+        conv_kernel_size: int = 31,
+        flash_attn: bool = False,
+    ):
+        super().__init__()
+        self.fc_factor = 0.5
+        self.norm_feed_forward1 = nn.LayerNorm(d_model)
+        self.feed_forward1 = ConformerFeedForward(d_model=d_model, d_ff=d_ff)
+        self.norm_conv = nn.LayerNorm(d_model)
+        self.conv = ConformerConvolution(
+            d_model=d_model,
+            kernel_size=conv_kernel_size,
+        )
+        self.norm_self_att = nn.LayerNorm(d_model)
+        if self_attention_model == "rotary":
+            self.self_attn: nn.Module = RotaryPositionMultiHeadAttention(
+                n_head=n_heads,
+                n_feat=d_model,
+                flash_attn=flash_attn,
+            )
+        else:
+            assert not flash_attn, "Not supported flash_attn for rel_pos"
+            self.self_attn = RelPositionMultiHeadAttention(
+                n_head=n_heads,
+                n_feat=d_model,
+            )
+        self.norm_feed_forward2 = nn.LayerNorm(d_model)
+        self.feed_forward2 = ConformerFeedForward(d_model=d_model, d_ff=d_ff)
+        self.norm_out = nn.LayerNorm(d_model)
+
+    def forward(
+        self,
+        x: Tensor,
+        pos_emb: Union[Tensor, List[Tensor]],
+        att_mask: Optional[Tensor] = None,
+        pad_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        residual = x
+        x = self.norm_feed_forward1(x)
+        x = self.feed_forward1(x)
+        residual = residual + x * self.fc_factor
+
+        x = self.norm_self_att(residual)
+        x = self.self_attn(x, x, x, pos_emb, mask=att_mask)
+        residual = residual + x
+
+        x = self.norm_conv(residual)
+        x = self.conv(x, pad_mask=pad_mask)
+        residual = residual + x
+
+        x = self.norm_feed_forward2(residual)
+        x = self.feed_forward2(x)
+        residual = residual + x * self.fc_factor
+
+        x = self.norm_out(residual)
+        return x
+
+
+class ConformerEncoder(nn.Module):
+    """
+    Conformer Encoder module.
+    This module encapsulates the entire Conformer encoder architecture,
+    consisting of a StridingSubsampling layer, positional embeddings, and
+    a stack of Conformer Layers.
+    It serves as the main component responsible for processing speech features.
+    """
+
+    def __init__(
+        self,
+        feat_in: int = 64,
+        n_layers: int = 16,
+        d_model: int = 768,
+        subsampling_factor: int = 4,
+        ff_expansion_factor: int = 4,
+        self_attention_model: str = "rotary",
+        n_heads: int = 16,
+        pos_emb_max_len: int = 5000,
+        conv_kernel_size: int = 31,
+        flash_attn: bool = False,
+    ):
+        super().__init__()
+        self.feat_in = feat_in
+        assert self_attention_model in [
+            "rotary",
+            "rel_pos",
+        ], f"Not supported attn = {self_attention_model}"
+
+        self.pre_encode = StridingSubsampling(
+            subsampling_factor=subsampling_factor,
+            feat_in=feat_in,
+            feat_out=d_model,
+            conv_channels=d_model,
+        )
+
+        if self_attention_model == "rotary":
+            self.pos_enc: nn.Module = RotaryPositionalEmbedding(
+                d_model // n_heads, pos_emb_max_len
+            )
+        else:
+            self.pos_enc = RelPositionalEmbedding(d_model, pos_emb_max_len)
+
+        self.layers = nn.ModuleList()
+        for _ in range(n_layers):
+            layer = ConformerLayer(
+                d_model=d_model,
+                d_ff=d_model * ff_expansion_factor,
+                self_attention_model=self_attention_model,
+                n_heads=n_heads,
+                conv_kernel_size=conv_kernel_size,
+                flash_attn=flash_attn,
+            )
+            self.layers.append(layer)
+
+        self.pos_enc.extend_pe(pos_emb_max_len, next(self.parameters()).device)
+
+    def input_example(
+        self,
+        batch_size: int = 1,
+        seqlen: int = 200,
+    ):
+        device = next(self.parameters()).device
+        features = torch.zeros(batch_size, self.feat_in, seqlen)
+        feature_lengths = torch.full([batch_size], features.shape[-1])
+        return features.float().to(device), feature_lengths.to(device)
+
+    def input_names(self):
+        return ["audio_signal", "length"]
+
+    def output_names(self):
+        return ["encoded", "encoded_len"]
+
+    def dynamic_axes(self):
+        return {
+            "audio_signal": {0: "batch_size", 2: "seq_len"},
+            "length": {0: "batch_size"},
+            "encoded": {0: "batch_size", 1: "seq_len"},
+            "encoded_len": {0: "batch_size"},
+        }
+
+    def forward(self, audio_signal: Tensor, length: Tensor) -> Tuple[Tensor, Tensor]:
+        audio_signal, length = self.pre_encode(
+            x=audio_signal.transpose(1, 2), lengths=length
+        )
+
+        max_len = audio_signal.size(1)
+        audio_signal, pos_emb = self.pos_enc(x=audio_signal)
+
+        pad_mask = torch.arange(0, max_len, device=audio_signal.device).expand(
+            length.size(0), -1
+        ) < length.unsqueeze(-1)
+
+        att_mask = None
+        if audio_signal.shape[0] > 1:
+            att_mask = pad_mask.unsqueeze(1).repeat([1, max_len, 1])
+            att_mask = torch.logical_and(att_mask, att_mask.transpose(1, 2))
+            att_mask = ~att_mask
+
+        pad_mask = ~pad_mask
+
+        for layer in self.layers:
+            audio_signal = layer(
+                x=audio_signal,
+                pos_emb=pos_emb,
+                att_mask=att_mask,
+                pad_mask=pad_mask,
+            )
+
+        return audio_signal.transpose(1, 2), length

gigaam_transformers.py

@@ -0,0 +1,447 @@
+from typing import Dict, List, Optional, Union, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio
+from .encoder import ConformerEncoder
+from torch import Tensor
+from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2Processor
+from transformers.configuration_utils import PretrainedConfig
+from transformers.feature_extraction_sequence_utils import \
+    SequenceFeatureExtractor
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.modeling_outputs import CausalLMOutput, Seq2SeqLMOutput
+from transformers.modeling_utils import PreTrainedModel
+
+
+class GigaAMCTC(nn.Module):
+    """
+    GigaAM-CTC model
+    """
+
+    def __init__(self, config_encoder, config_head):
+        super().__init__()
+        self.encoder = ConformerEncoder(**config_encoder)
+        self.head = CTCHead(**config_head)
+
+    def forward(self, input_features: Tensor, input_lengths: Tensor) -> Tensor:
+        encoded, encoded_lengths = self.encoder(input_features, input_lengths)
+        logits = self.head(encoded)
+        return logits, encoded_lengths
+
+
+class GigaAMRNNT(nn.Module):
+    """
+    GigaAM-RNNT model
+    """
+
+    def __init__(self, config_encoder, config_head):
+        super().__init__()
+        self.encoder = ConformerEncoder(**config_encoder)
+        self.head = RNNTHead(**config_head)
+
+    def forward(self, input_features: Tensor, input_lengths: Tensor, targets: Tensor, target_lengths: Tensor) -> Tensor:
+        encoded, encoded_lengths = self.encoder(input_features, input_lengths)
+        # During training, loss must be computed, so decoder forward is necessary
+        decoder_out, target_lengths, states = self.head.decoder(targets=targets, target_length=target_lengths)
+        joint = self.head.joint(encoder_outputs=encoded, decoder_outputs=decoder_out)
+        # loss = self.loss(
+        #     log_probs=joint, targets=targets, input_lengths=encoded_lengths, target_lengths=target_lengths
+        # )
+        
+        return joint, encoded_lengths
+
+
+class CTCHead(nn.Module):
+    """
+    CTC Head module for Connectionist Temporal Classification.
+    """
+
+    def __init__(self, feat_in: int, num_classes: int):
+        super().__init__()
+        self.decoder_layers = nn.Sequential(
+            nn.Conv1d(feat_in, num_classes, kernel_size=1)
+        )
+
+    def forward(self, encoder_output: Tensor) -> Tensor:
+        # B x C x T
+        return self.decoder_layers(encoder_output)
+
+
+class RNNTJoint(nn.Module):
+    """
+    RNN-Transducer Joint Network Module.
+    This module combines the outputs of the encoder and the prediction network using
+    a linear transformation followed by ReLU activation and another linear projection.
+    """
+
+    def __init__(
+        self, enc_hidden: int, pred_hidden: int, joint_hidden: int, num_classes: int
+    ):
+        super().__init__()
+        self.enc_hidden = enc_hidden
+        self.pred_hidden = pred_hidden
+        self.pred = nn.Linear(pred_hidden, joint_hidden)
+        self.enc = nn.Linear(enc_hidden, joint_hidden)
+        self.joint_net = nn.Sequential(nn.ReLU(), nn.Linear(joint_hidden, num_classes))
+
+    def joint(self, encoder_out: Tensor, decoder_out: Tensor) -> Tensor:
+        """
+        Combine the encoder and prediction network outputs into a joint representation.
+        """
+        enc = self.enc(encoder_out).unsqueeze(2)
+        pred = self.pred(decoder_out).unsqueeze(1)
+        return self.joint_net(enc + pred)
+
+    def input_example(self):
+        device = next(self.parameters()).device
+        enc = torch.zeros(1, self.enc_hidden, 1)
+        dec = torch.zeros(1, self.pred_hidden, 1)
+        return enc.float().to(device), dec.float().to(device)
+
+    def input_names(self):
+        return ["enc", "dec"]
+
+    def output_names(self):
+        return ["joint"]
+
+    def forward(self, enc: Tensor, dec: Tensor) -> Tensor:
+        return self.joint(enc.transpose(1, 2), dec.transpose(1, 2))
+
+
+class RNNTDecoder(nn.Module):
+    """
+    RNN-Transducer Decoder Module.
+    This module handles the prediction network part of the RNN-Transducer architecture.
+    """
+
+    def __init__(self, pred_hidden: int, pred_rnn_layers: int, num_classes: int):
+        super().__init__()
+        self.blank_id = num_classes - 1
+        self.pred_hidden = pred_hidden
+        self.embed = nn.Embedding(num_classes, pred_hidden, padding_idx=self.blank_id)
+        self.lstm = nn.LSTM(pred_hidden, pred_hidden, pred_rnn_layers)
+
+    def predict(
+        self,
+        x: Optional[Tensor],
+        state: Optional[Tensor],
+        batch_size: int = 1,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Make predictions based on the current input and previous states.
+        If no input is provided, use zeros as the initial input.
+        """
+        if x is not None:
+            emb: Tensor = self.embed(x)
+        else:
+            emb = torch.zeros(
+                (batch_size, 1, self.pred_hidden), device=next(self.parameters()).device
+            )
+        g, hid = self.lstm(emb.transpose(0, 1), state)
+        return g.transpose(0, 1), hid
+
+    def input_example(self):
+        device = next(self.parameters()).device
+        label = torch.tensor([[0]]).to(device)
+        hidden_h = torch.zeros(1, 1, self.pred_hidden).to(device)
+        hidden_c = torch.zeros(1, 1, self.pred_hidden).to(device)
+        return label, hidden_h, hidden_c
+
+    def input_names(self):
+        return ["x", "h", "c"]
+
+    def output_names(self):
+        return ["dec", "h", "c"]
+
+    def forward(self, x: Tensor, h: Tensor, c: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        ONNX-specific forward with x, state = (h, c) -> x, h, c.
+        """
+        emb = self.embed(x)
+        g, (h, c) = self.lstm(emb.transpose(0, 1), (h, c))
+        return g.transpose(0, 1), h, c
+
+
+class RNNTHead(nn.Module):
+    """
+    RNN-Transducer Head Module.
+    This module combines the decoder and joint network components of the RNN-Transducer architecture.
+    """
+
+    def __init__(self, decoder: Dict[str, int], joint: Dict[str, int]):
+        super().__init__()
+        self.decoder = RNNTDecoder(**decoder)
+        self.joint = RNNTJoint(**joint)
+
+
+class GigaAMFeatureExtractor(SequenceFeatureExtractor):
+    """
+    Feature extractor for GigaAM.
+    """
+    model_input_names = ["input_features"]
+
+    def __init__(
+        self,
+        feature_size=64,
+        sampling_rate=16000,
+        padding_value=0.0,
+        chunk_length=30.0,
+        **kwargs,
+    ):
+        super().__init__(
+            feature_size=feature_size,
+            sampling_rate=sampling_rate,
+            padding_value=padding_value,
+            chunk_length=chunk_length,
+            **kwargs,
+        )
+        self.hop_length = sampling_rate // 100
+        self.n_samples = chunk_length * sampling_rate
+        self.featurizer = torchaudio.transforms.MelSpectrogram(
+                sample_rate=sampling_rate,
+                n_fft=sampling_rate // 40,
+                win_length=sampling_rate // 40,
+                hop_length=self.hop_length,
+                n_mels=feature_size,
+            )
+
+    def to_dict(self) -> Dict[str, Union[str, int, Dict]]:
+        dictionary = super().to_dict()
+
+        if "featurizer" in dictionary:
+            del dictionary["featurizer"]
+        dictionary["hop_length"] = self.hop_length
+        dictionary["n_samples"] = self.n_samples
+        return dictionary
+
+    def out_len(self, input_lengths: Tensor) -> Tensor:
+        """
+        Calculates the output length after the feature extraction process.
+        """
+        return input_lengths.div(self.hop_length, rounding_mode="floor").add(1).long()
+
+    def __call__(
+        self,
+        raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
+        sampling_rate: Optional[int] = None,
+        padding: str = "max_length",
+        **kwargs,
+    ):
+        is_batched_numpy = (
+            isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
+        )
+        if is_batched_numpy and len(raw_speech.shape) > 2:
+            raise ValueError(
+                f"Only mono-channel audio is supported for input to {self}"
+            )
+        is_batched = is_batched_numpy or (
+            isinstance(raw_speech, (list, tuple))
+            and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
+        )
+
+        if is_batched:
+            raw_speech = [
+                np.asarray([speech], dtype=np.float32).T for speech in raw_speech
+            ]
+        elif not is_batched and not isinstance(raw_speech, np.ndarray):
+            raw_speech = np.asarray(raw_speech, dtype=np.float32)
+        elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(
+            np.float64
+        ):
+            raw_speech = raw_speech.astype(np.float32)
+
+        # always return batch
+        if not is_batched:
+            raw_speech = [np.asarray([raw_speech]).T]
+
+        input_lengths = torch.tensor([len(speech) for speech in raw_speech])
+
+        batched_speech = BatchFeature({"input_features": raw_speech})
+
+        padded_inputs = self.pad(
+            batched_speech,
+            padding=padding,
+            max_length=self.n_samples,
+            truncation=False,
+            return_tensors="pt",
+        )
+
+        input_features = padded_inputs["input_features"].transpose(1, 2)
+        input_features = self.featurizer(input_features).squeeze(1)
+        input_features = torch.log(input_features.clamp_(1e-9, 1e9))
+        input_lengths = self.out_len(input_lengths)
+
+        return BatchFeature({"input_features": input_features, "input_lengths": input_lengths}, tensor_type="pt")
+
+
+class GigaAMCTCTokenizer(Wav2Vec2CTCTokenizer):
+    """
+    Char tokenizer for GigaAM-CTC model.
+    """
+    def __init__(
+        self,
+        vocab_file,
+        unk_token="[BLANK]",
+        pad_token="[BLANK]",
+        bos_token=None,
+        eos_token=None,
+        word_delimiter_token=" ",
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            word_delimiter_token=word_delimiter_token,
+            **kwargs,
+        )
+
+
+class GigaAMProcessor(Wav2Vec2Processor):
+    feature_extractor_class = "GigaAMFeatureExtractor"
+    tokenizer_class = "GigaAMCTCTokenizer"
+
+    def __init__(self, feature_extractor, tokenizer):
+        # super().__init__(feature_extractor, tokenizer)
+        self.feature_extractor = feature_extractor
+        self.tokenizer = tokenizer
+        self.current_processor = self.feature_extractor
+        self._in_target_context_manager = False
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        feature_extractor = GigaAMFeatureExtractor.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        tokenizer = GigaAMCTCTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        return cls(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+
+class GigaAMConfig(PretrainedConfig):
+    model_type = "gigaam-ctc"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+class GigaAMCTCHF(PreTrainedModel):
+    """
+    GigaAM-CTC model for transformers
+    """
+    config_class = GigaAMConfig
+    base_model_prefix = "gigaamctc"
+    main_input_name = "input_features"
+
+    def __init__(self, config: GigaAMConfig):
+        super().__init__(config)
+        self.model = GigaAMCTC(config.encoder, config.head)
+
+    def forward(self, input_features, input_lengths, labels=None, **kwargs):
+
+        # B x C x T
+        logits, encoded_lengths = self.model(input_features, input_lengths)
+        # B x C x T -> B x T x C -> T x B x C
+        log_probs = torch.log_softmax(
+            logits.transpose(1, 2), dim=-1, dtype=torch.float32
+        ).transpose(0, 1)
+
+        loss = None
+        if labels is not None:
+            labels_mask = labels >= 0
+            target_lengths = labels_mask.sum(-1)
+            flattened_targets = labels.masked_select(labels_mask)
+
+            loss = nn.functional.ctc_loss(
+                log_probs,
+                flattened_targets,
+                encoded_lengths,
+                target_lengths,
+                blank=self.config.blank_id,
+                zero_infinity=True,
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits.transpose(1, 2))
+
+
+class GigaAMRNNTHF(PreTrainedModel):
+    """
+    GigaAM-RNNT model for transformers
+    """
+    config_class = GigaAMConfig
+    base_model_prefix = "gigaamrnnt"
+    main_input_name = "input_features"
+
+    def __init__(self, config: GigaAMConfig):
+        super().__init__(config)
+        self.model = GigaAMRNNT(config.encoder, config.head)
+
+    def forward(self, input_features, input_lengths, labels=None, **kwargs):
+
+        # B x C x T
+        encoder_out, encoded_lengths = self.model.encoder(input_features, input_lengths)
+        encoder_out = encoder_out.transpose(1, 2)
+        batch_size = encoder_out.shape[0]
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(torch.int32)
+            labels_mask = labels >= 0
+            target_lengths = labels_mask.sum(-1).to(torch.int32)
+
+            hidden_states = torch.zeros((self.config.head["decoder"]["pred_rnn_layers"], batch_size, self.model.head.decoder.pred_hidden), device=encoder_out.device)
+            hidden_c = torch.zeros((self.config.head["decoder"]["pred_rnn_layers"], batch_size, self.model.head.decoder.pred_hidden), device=encoder_out.device)
+            plus_one_dim = self.config.blank_id * torch.ones((batch_size, 1), dtype=torch.int32, device=encoder_out.device)
+            
+            decoder_out, h, c = self.model.head.decoder(torch.cat((plus_one_dim, labels), dim=1), hidden_states, hidden_c)
+
+            joint = self.model.head.joint.joint(encoder_out, decoder_out)
+            loss = torchaudio.functional.rnnt_loss(
+                logits=joint,
+                targets=labels,
+                logit_lengths=encoded_lengths,
+                target_lengths=target_lengths,
+                blank=self.config.blank_id,
+            )
+
+        return Seq2SeqLMOutput(loss=loss, logits=encoder_out.transpose(1, 2))
+    
+    def _greedy_decode(self, x: Tensor, seqlen: Tensor) -> str:
+        """
+        Internal helper function for performing greedy decoding on a single sequence.
+        """
+        hyp: List[int] = []
+        dec_state: Optional[Tensor] = None
+        last_label: Optional[Tensor] = None
+        for t in range(seqlen):
+            f = x[t, :, :].unsqueeze(1)
+            not_blank = True
+            new_symbols = 0
+            while not_blank and new_symbols < self.config.max_symbols:
+                g, hidden = self.model.head.decoder.predict(last_label, dec_state)
+                k = self.model.head.joint.joint(f, g)[0, 0, 0, :].argmax(0).item()
+                if k == self.config.blank_id:
+                    not_blank = False
+                else:
+                    hyp.append(k)
+                    dec_state = hidden
+                    last_label = torch.tensor([[hyp[-1]]]).to(x.device)
+                    new_symbols += 1
+
+        return torch.tensor([hyp], dtype=torch.int32)
+
+    @torch.inference_mode()
+    def generate(self, input_features: Tensor, input_lengths: Tensor, **kwargs) -> torch.Tensor:
+        """
+        Decode the output of an RNN-T model into a list of hypotheses.
+        """
+        encoder_out, encoded_lengths = self.model.encoder(input_features, input_lengths)
+        encoder_out = encoder_out.transpose(1, 2)
+        b = encoder_out.shape[0]
+        preds = []
+        for i in range(b):
+            inseq = encoder_out[i, :, :].unsqueeze(1)
+            preds.append(self._greedy_decode(inseq, encoded_lengths[i]))
+        return torch.cat(preds, dim=1)

preprocessor_config.json

@@ -0,0 +1,18 @@
+{
+    "chunk_length": 30,
+    "feature_extractor_type": "GigaAMFeatureExtractor",
+    "feature_extractor_class": "GigaAMFeatureExtractor",
+    "feature_size": 64,
+    "hop_length": 160,
+    "n_samples": 480000,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "return_attention_mask": true,
+    "sampling_rate": 16000,
+    "auto_map": {
+      "AutoFeatureExtractor": "gigaam_transformers.GigaAMFeatureExtractor",
+      "AutoProcessor": "gigaam_transformers.GigaAMProcessor"
+    },
+    "processor_class": "GigaAMProcessor",
+    "model_type": "gigaam-rnnt"
+  }

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16888332b279f5296b411f5bbe5385f20448d261b2c9a6425e1a83026efb2018
+size 935306114

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+    "pad_token": "[BLANK]",
+    "unk_token": "[BLANK]"
+  }
+  

tokenizer_config.json

@@ -0,0 +1,27 @@
+{
+    "added_tokens_decoder": {
+      "33": {
+        "content": "[BLANK]",
+        "lstrip": true,
+        "normalized": false,
+        "rstrip": true,
+        "single_word": false,
+        "special": false
+      }
+    },
+    "bos_token": null,
+    "clean_up_tokenization_spaces": false,
+    "do_lower_case": false,
+    "eos_token": null,
+    "model_max_length": 1000,
+    "pad_token": "[BLANK]",
+    "replace_word_delimiter_char": " ",
+    "target_lang": null,
+    "tokenizer_class": "GigaAMTokenizer",
+    "unk_token": "[BLANK]",
+    "word_delimiter_token": " ",
+    "auto_map": {
+      "AutoTokenizer": ["gigaam_transformers.GigaAMTokenizer", null]
+    }
+  }
+  

vocab.json

@@ -0,0 +1,37 @@
+{
+    " ": 0,
+    "[BLANK]": 33,
+    "а": 1,
+    "б": 2,
+    "в": 3,
+    "г": 4,
+    "д": 5,
+    "е": 6,
+    "ж": 7,
+    "з": 8,
+    "и": 9,
+    "й": 10,
+    "к": 11,
+    "л": 12,
+    "м": 13,
+    "н": 14,
+    "о": 15,
+    "п": 16,
+    "р": 17,
+    "с": 18,
+    "т": 19,
+    "у": 20,
+    "ф": 21,
+    "х": 22,
+    "ц": 23,
+    "ч": 24,
+    "ш": 25,
+    "щ": 26,
+    "ъ": 27,
+    "ы": 28,
+    "ь": 29,
+    "э": 30,
+    "ю": 31,
+    "я": 32
+  }
+