funasr sdk

游雁

2023-07-05 4e2fe544ae37174a3e09dfcdbbdae5abfe711e53

 funasr/models/e2e_diar_eend_ola.py

@@ -9,14 +9,13 @@
import numpy as np
import torch
import torch.nn as  nn
from typeguard import check_argument_types

from funasr.models.frontend.wav_frontend import WavFrontendMel23
from funasr.modules.eend_ola.encoder import EENDOLATransformerEncoder
from funasr.modules.eend_ola.encoder_decoder_attractor import EncoderDecoderAttractor
from funasr.modules.eend_ola.utils.power import generate_mapping_dict
from funasr.torch_utils.device_funcs import force_gatherable
from funasr.train.abs_espnet_model import AbsESPnetModel
from funasr.models.base_model import FunASRModel

if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
    pass
@@ -34,7 +33,7 @@
    return att


class DiarEENDOLAModel(AbsESPnetModel):
class DiarEENDOLAModel(FunASRModel):
    """EEND-OLA diarization model"""

    def __init__(
@@ -48,7 +47,6 @@
            mapping_dict=None,
            **kwargs,
    ):
        assert check_argument_types()

        super().__init__()
        self.frontend = frontend
@@ -91,7 +89,6 @@
            text_lengths: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
        """Frontend + Encoder + Decoder + Calc loss

        Args:
            speech: (Batch, Length, ...)
            speech_lengths: (Batch, )