freedakgmail
6116b1f3c6
Update API Documentation / build-api-docs (push) Has been cancelled
Add complete FunASR codebase including models, runtime, and documentation.
31 lines
811 B
Python
31 lines
811 B
Python
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):
|
|
"""Initialize SimpleAvg.
|
|
|
|
Args:
|
|
feat_dim: Size/dimension parameter.
|
|
"""
|
|
super(SimpleAvg, self).__init__()
|
|
self.feat_dim = feat_dim
|
|
|
|
def forward(self, x, ilens):
|
|
"""Forward pass for training.
|
|
|
|
Args:
|
|
x: TODO.
|
|
ilens: TODO.
|
|
"""
|
|
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:
|
|
"""Output size."""
|
|
return self.feat_dim
|