Upload ConformerEncoder

config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "ConformerEncoder"
+  ],
+  "auto_map": {
+    "AutoConfig": "conformer.ConformerConfig",
+    "AutoModel": "conformer.ConformerEncoder"
+  },
+  "conformer_depthwise_conv_kernel_size": 31,
+  "conformer_dropout": 0.1,
+  "conformer_ffn_dim": 576,
+  "conformer_input_dim": 144,
+  "conformer_num_heads": 4,
+  "conformer_num_layers": 8,
+  "input_dim": 80,
+  "model_type": "conformer",
+  "output_dim": 40,
+  "time_reduction_stride": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.35.2"
+}

conformer.py

@@ -0,0 +1,40 @@
+from torchaudio.models import Conformer
+from torchaudio.models.rnnt import _TimeReduction
+from transformers import PretrainedConfig, PreTrainedModel
+import torch
+from typing import List, Tuple, Optional
+
+
+class ConformerConfig(PretrainedConfig):
+    model_type = 'conformer'
+
+
+class ConformerEncoder(PreTrainedModel):
+    def __init__(
+        self,
+        config,
+    ) -> None:
+        super().__init__(config)
+        self.time_reduction = _TimeReduction(config.time_reduction_stride)
+        self.input_linear = torch.nn.Linear(
+            config.input_dim * config.time_reduction_stride,
+            config.conformer_input_dim)
+        self.conformer = Conformer(
+            num_layers=config.conformer_num_layers,
+            input_dim=config.conformer_input_dim,
+            ffn_dim=config.conformer_ffn_dim,
+            num_heads=config.conformer_num_heads,
+            depthwise_conv_kernel_size=config.conformer_depthwise_conv_kernel_size,
+            dropout=config.conformer_dropout,
+            use_group_norm=True,
+            convolution_first=True,
+        )
+        self.output_linear = torch.nn.Linear(config.conformer_input_dim, config.output_dim)
+
+    def forward(self, input: torch.Tensor,
+                lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        time_reduction_out, time_reduction_lengths = self.time_reduction(input, lengths)
+        input_linear_out = self.input_linear(time_reduction_out)
+        x, lengths = self.conformer(input_linear_out, time_reduction_lengths)
+        output_linear_out = self.output_linear(x)
+        return output_linear_out, lengths

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98b344ee3d7122a46057f410108a03413a6bb69cbe78496c4a3c35a50c5fc1c0
+size 15780592