first commit for takway.ai

funasr/models/whisper_lid/eres2net/simple_avg.py

@@ -0,0 +1,18 @@
+import torch
+
+from funasr.models.encoder.abs_encoder import AbsEncoder
+from funasr.modules.nets_utils import make_pad_mask
+
+
+class SimpleAvg(AbsEncoder):
+    def __init__(self, feat_dim):
+        super(SimpleAvg, self).__init__()
+        self.feat_dim = feat_dim
+
+    def forward(self, x, ilens):
+        mask = ~make_pad_mask(ilens, maxlen=x.shape[1]).to(x.device)
+        avg_x = (x * mask[:, :, None]).sum(1) / mask.sum(-1)[:, None]
+        return avg_x
+
+    def output_size(self) -> int:
+        return self.feat_dim