Upload ConformerEncoder

config.json

@@ -12,9 +12,12 @@
   "conformer_input_dim": 144,
   "conformer_num_heads": 4,
   "conformer_num_layers": 8,
+  "ctc_loss_reduction": "mean",
+  "ctc_zero_infinity": true,
   "input_dim": 80,
   "model_type": "conformer",
   "output_dim": 40,
+  "pad_token_id": 39,
   "time_reduction_stride": 4,
   "torch_dtype": "float32",
   "transformers_version": "4.35.2"

conformer.py

@@ -2,6 +2,7 @@ from torchaudio.models import Conformer
 from torchaudio.models.rnnt import _TimeReduction
 from transformers import PretrainedConfig, PreTrainedModel
 import torch
+from torch import nn
 from typing import List, Tuple, Optional
 
 
@@ -33,10 +34,33 @@ class ConformerEncoder(PreTrainedModel):
         )
         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)
+    def forward(self, inputs, lengths, labels=None):
+        time_reduction_out, time_reduction_lengths = self.time_reduction(inputs, 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
+        x, input_lengths = self.conformer(input_linear_out, time_reduction_lengths)
+        logits = self.output_linear(x)
+
+        loss = None
+        if labels is not None:
+            labels_mask = labels >= 0
+            target_lengths = labels_mask.sum(-1)
+            flattened_targets = labels.masked_select(labels_mask)
+            log_probs = nn.functional.log_softmax(
+                logits,
+                dim=-1,
+                dtype=torch.float32
+            ).transpose(0, 1)
+
+            with torch.backends.cudnn.flags(enabled=False):
+                loss = nn.functional.ctc_loss(
+                    log_probs,
+                    flattened_targets,
+                    input_lengths,
+                    target_lengths,
+                    blank=self.config.pad_token_id,
+                    reduction=self.config.ctc_loss_reduction,
+                    zero_infinity=self.config.ctc_zero_infinity,
+                )
+
+        output = (logits, input_lengths)
+        return ((loss,) + output) if loss is not None else output

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c48c1cb2f04deb1086872dfa916773384f36d7a755c2ee220170290851ddfd46
+oid sha256:4750d570d1888762e0c5c89883addd1ef8914ff0a46b9c85ea931c982f85285a
 size 15780592