Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/models/contextual_paraformer/decoder.py

@@ -0,0 +1,507 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+import logging
+import numpy as np
+from typing import Tuple
+
+from funasr.register import tables
+from funasr.models.scama import utils as myutils
+from funasr.models.transformer.utils.repeat import repeat
+from funasr.models.transformer.layer_norm import LayerNorm
+from funasr.models.transformer.embedding import PositionalEncoding
+from funasr.models.paraformer.decoder import DecoderLayerSANM, ParaformerSANMDecoder
+from funasr.models.sanm.positionwise_feed_forward import PositionwiseFeedForwardDecoderSANM
+from funasr.models.sanm.attention import (
+    MultiHeadedAttentionSANMDecoder,
+    MultiHeadedAttentionCrossAtt,
+)
+
+
+class ContextualDecoderLayer(torch.nn.Module):
+    def __init__(
+        self,
+        size,
+        self_attn,
+        src_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an DecoderLayer object."""
+        super(ContextualDecoderLayer, self).__init__()
+        self.size = size
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(size)
+        if self_attn is not None:
+            self.norm2 = LayerNorm(size)
+        if src_attn is not None:
+            self.norm3 = LayerNorm(size)
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear1 = torch.nn.Linear(size + size, size)
+            self.concat_linear2 = torch.nn.Linear(size + size, size)
+
+    def forward(
+        self,
+        tgt,
+        tgt_mask,
+        memory,
+        memory_mask,
+        cache=None,
+    ):
+        # tgt = self.dropout(tgt)
+        """Forward pass for training.
+        
+            Args:
+                tgt: TODO.
+                tgt_mask: TODO.
+                memory: TODO.
+                memory_mask: TODO.
+                cache: State cache dict for streaming inference.
+            """
+        if isinstance(tgt, Tuple):
+            tgt, _ = tgt
+        residual = tgt
+        if self.normalize_before:
+            tgt = self.norm1(tgt)
+        tgt = self.feed_forward(tgt)
+
+        x = tgt
+        if self.normalize_before:
+            tgt = self.norm2(tgt)
+        if self.training:
+            cache = None
+        x, cache = self.self_attn(tgt, tgt_mask, cache=cache)
+        x = residual + self.dropout(x)
+        x_self_attn = x
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm3(x)
+        x = self.src_attn(x, memory, memory_mask)
+        x_src_attn = x
+
+        x = residual + self.dropout(x)
+        return x, tgt_mask, x_self_attn, x_src_attn
+
+
+class ContextualBiasDecoder(torch.nn.Module):
+    def __init__(
+        self,
+        size,
+        src_attn,
+        dropout_rate,
+        normalize_before=True,
+    ):
+        """Construct an DecoderLayer object."""
+        super(ContextualBiasDecoder, self).__init__()
+        self.size = size
+        self.src_attn = src_attn
+        if src_attn is not None:
+            self.norm3 = LayerNorm(size)
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.normalize_before = normalize_before
+
+    def forward(self, tgt, tgt_mask, memory, memory_mask=None, cache=None):
+        """Forward pass for training.
+        
+            Args:
+                tgt: TODO.
+                tgt_mask: TODO.
+                memory: TODO.
+                memory_mask: TODO.
+                cache: State cache dict for streaming inference.
+            """
+        x = tgt
+        if self.src_attn is not None:
+            if self.normalize_before:
+                x = self.norm3(x)
+            x = self.dropout(self.src_attn(x, memory, memory_mask))
+        return x, tgt_mask, memory, memory_mask, cache
+
+
+@tables.register("decoder_classes", "ContextualParaformerDecoder")
+class ContextualParaformerDecoder(ParaformerSANMDecoder):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
+    https://arxiv.org/abs/2006.01713
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder_output_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        self_attention_dropout_rate: float = 0.0,
+        src_attention_dropout_rate: float = 0.0,
+        input_layer: str = "embed",
+        use_output_layer: bool = True,
+        pos_enc_class=PositionalEncoding,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        att_layer_num: int = 6,
+        kernel_size: int = 21,
+        sanm_shfit: int = 0,
+    ):
+        """Initialize ContextualParaformerDecoder.
+        
+            Args:
+                vocab_size: Size/dimension parameter.
+                encoder_output_size: Size/dimension parameter.
+                attention_heads: TODO.
+                linear_units: TODO.
+                num_blocks: TODO.
+                dropout_rate: TODO.
+                positional_dropout_rate: TODO.
+                self_attention_dropout_rate: TODO.
+                src_attention_dropout_rate: TODO.
+                input_layer: TODO.
+                use_output_layer: TODO.
+                pos_enc_class: TODO.
+                normalize_before: TODO.
+                concat_after: TODO.
+                att_layer_num: TODO.
+                kernel_size: Size/dimension parameter.
+                sanm_shfit: TODO.
+            """
+        super().__init__(
+            vocab_size=vocab_size,
+            encoder_output_size=encoder_output_size,
+            dropout_rate=dropout_rate,
+            positional_dropout_rate=positional_dropout_rate,
+            input_layer=input_layer,
+            use_output_layer=use_output_layer,
+            pos_enc_class=pos_enc_class,
+            normalize_before=normalize_before,
+        )
+
+        attention_dim = encoder_output_size
+        if input_layer == "none":
+            self.embed = None
+        if input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(vocab_size, attention_dim),
+                # pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(vocab_size, attention_dim),
+                torch.nn.LayerNorm(attention_dim),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        else:
+            raise ValueError(f"only 'embed' or 'linear' is supported: {input_layer}")
+
+        self.normalize_before = normalize_before
+        if self.normalize_before:
+            self.after_norm = LayerNorm(attention_dim)
+        if use_output_layer:
+            self.output_layer = torch.nn.Linear(attention_dim, vocab_size)
+        else:
+            self.output_layer = None
+
+        self.att_layer_num = att_layer_num
+        self.num_blocks = num_blocks
+        if sanm_shfit is None:
+            sanm_shfit = (kernel_size - 1) // 2
+        self.decoders = repeat(
+            att_layer_num - 1,
+            lambda lnum: DecoderLayerSANM(
+                attention_dim,
+                MultiHeadedAttentionSANMDecoder(
+                    attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit
+                ),
+                MultiHeadedAttentionCrossAtt(
+                    attention_heads, attention_dim, src_attention_dropout_rate
+                ),
+                PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.bias_decoder = ContextualBiasDecoder(
+            size=attention_dim,
+            src_attn=MultiHeadedAttentionCrossAtt(
+                attention_heads, attention_dim, src_attention_dropout_rate
+            ),
+            dropout_rate=dropout_rate,
+            normalize_before=True,
+        )
+        self.bias_output = torch.nn.Conv1d(attention_dim * 2, attention_dim, 1, bias=False)
+        self.last_decoder = ContextualDecoderLayer(
+            attention_dim,
+            MultiHeadedAttentionSANMDecoder(
+                attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=sanm_shfit
+            ),
+            MultiHeadedAttentionCrossAtt(
+                attention_heads, attention_dim, src_attention_dropout_rate
+            ),
+            PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+            dropout_rate,
+            normalize_before,
+            concat_after,
+        )
+        if num_blocks - att_layer_num <= 0:
+            self.decoders2 = None
+        else:
+            self.decoders2 = repeat(
+                num_blocks - att_layer_num,
+                lambda lnum: DecoderLayerSANM(
+                    attention_dim,
+                    MultiHeadedAttentionSANMDecoder(
+                        attention_dim, self_attention_dropout_rate, kernel_size, sanm_shfit=0
+                    ),
+                    None,
+                    PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                    dropout_rate,
+                    normalize_before,
+                    concat_after,
+                ),
+            )
+
+        self.decoders3 = repeat(
+            1,
+            lambda lnum: DecoderLayerSANM(
+                attention_dim,
+                None,
+                None,
+                PositionwiseFeedForwardDecoderSANM(attention_dim, linear_units, dropout_rate),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+
+    def forward(
+        self,
+        hs_pad: torch.Tensor,
+        hlens: torch.Tensor,
+        ys_in_pad: torch.Tensor,
+        ys_in_lens: torch.Tensor,
+        contextual_info: torch.Tensor,
+        clas_scale: float = 1.0,
+        return_hidden: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward decoder.
+
+        Args:
+            hs_pad: encoded memory, float32  (batch, maxlen_in, feat)
+            hlens: (batch)
+            ys_in_pad:
+                input token ids, int64 (batch, maxlen_out)
+                if input_layer == "embed"
+                input tensor (batch, maxlen_out, #mels) in the other cases
+            ys_in_lens: (batch)
+        Returns:
+            (tuple): tuple containing:
+
+            x: decoded token score before softmax (batch, maxlen_out, token)
+                if use_output_layer is True,
+            olens: (batch, )
+        """
+        tgt = ys_in_pad
+        tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None]
+
+        memory = hs_pad
+        memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :]
+
+        x = tgt
+        x, tgt_mask, memory, memory_mask, _ = self.decoders(x, tgt_mask, memory, memory_mask)
+        _, _, x_self_attn, x_src_attn = self.last_decoder(x, tgt_mask, memory, memory_mask)
+
+        # contextual paraformer related
+        contextual_length = torch.Tensor([contextual_info.shape[1]]).int().repeat(hs_pad.shape[0])
+        contextual_mask = myutils.sequence_mask(contextual_length, device=memory.device)[:, None, :]
+        cx, tgt_mask, _, _, _ = self.bias_decoder(
+            x_self_attn, tgt_mask, contextual_info, memory_mask=contextual_mask
+        )
+
+        if self.bias_output is not None:
+            x = torch.cat([x_src_attn, cx * clas_scale], dim=2)
+            x = self.bias_output(x.transpose(1, 2)).transpose(1, 2)  # 2D -> D
+            x = x_self_attn + self.dropout(x)
+
+        if self.decoders2 is not None:
+            x, tgt_mask, memory, memory_mask, _ = self.decoders2(x, tgt_mask, memory, memory_mask)
+
+        x, tgt_mask, memory, memory_mask, _ = self.decoders3(x, tgt_mask, memory, memory_mask)
+        if self.normalize_before:
+            x = self.after_norm(x)
+        olens = tgt_mask.sum(1)
+        if self.output_layer is not None and return_hidden is False:
+            x = self.output_layer(x)
+        return x, olens
+
+
+@tables.register("decoder_classes", "ContextualParaformerDecoderExport")
+class ContextualParaformerDecoderExport(torch.nn.Module):
+    def __init__(
+        self,
+        model,
+        max_seq_len=512,
+        model_name="decoder",
+        onnx: bool = True,
+        **kwargs,
+    ):
+        """Initialize ContextualParaformerDecoderExport.
+        
+            Args:
+                model: Model instance or model name.
+                max_seq_len: TODO.
+                model_name: TODO.
+                onnx: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+        from funasr.utils.torch_function import sequence_mask
+
+        self.model = model
+        self.make_pad_mask = sequence_mask(max_seq_len, flip=False)
+
+        from funasr.models.sanm.attention import MultiHeadedAttentionSANMDecoderExport
+        from funasr.models.sanm.attention import MultiHeadedAttentionCrossAttExport
+        from funasr.models.paraformer.decoder import DecoderLayerSANMExport
+        from funasr.models.transformer.positionwise_feed_forward import (
+            PositionwiseFeedForwardDecoderSANMExport,
+        )
+
+        for i, d in enumerate(self.model.decoders):
+            if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM):
+                d.feed_forward = PositionwiseFeedForwardDecoderSANMExport(d.feed_forward)
+            if isinstance(d.self_attn, MultiHeadedAttentionSANMDecoder):
+                d.self_attn = MultiHeadedAttentionSANMDecoderExport(d.self_attn)
+            if isinstance(d.src_attn, MultiHeadedAttentionCrossAtt):
+                d.src_attn = MultiHeadedAttentionCrossAttExport(d.src_attn)
+            self.model.decoders[i] = DecoderLayerSANMExport(d)
+
+        if self.model.decoders2 is not None:
+            for i, d in enumerate(self.model.decoders2):
+                if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM):
+                    d.feed_forward = PositionwiseFeedForwardDecoderSANMExport(d.feed_forward)
+                if isinstance(d.self_attn, MultiHeadedAttentionSANMDecoder):
+                    d.self_attn = MultiHeadedAttentionSANMDecoderExport(d.self_attn)
+                self.model.decoders2[i] = DecoderLayerSANMExport(d)
+
+        for i, d in enumerate(self.model.decoders3):
+            if isinstance(d.feed_forward, PositionwiseFeedForwardDecoderSANM):
+                d.feed_forward = PositionwiseFeedForwardDecoderSANMExport(d.feed_forward)
+            self.model.decoders3[i] = DecoderLayerSANMExport(d)
+
+        self.output_layer = model.output_layer
+        self.after_norm = model.after_norm
+        self.model_name = model_name
+
+        # bias decoder
+        if isinstance(self.model.bias_decoder.src_attn, MultiHeadedAttentionCrossAtt):
+            self.model.bias_decoder.src_attn = MultiHeadedAttentionCrossAttExport(
+                self.model.bias_decoder.src_attn
+            )
+        self.bias_decoder = self.model.bias_decoder
+
+        # last decoder
+        if isinstance(self.model.last_decoder.src_attn, MultiHeadedAttentionCrossAtt):
+            self.model.last_decoder.src_attn = MultiHeadedAttentionCrossAttExport(
+                self.model.last_decoder.src_attn
+            )
+        if isinstance(self.model.last_decoder.self_attn, MultiHeadedAttentionSANMDecoder):
+            self.model.last_decoder.self_attn = MultiHeadedAttentionSANMDecoderExport(
+                self.model.last_decoder.self_attn
+            )
+        if isinstance(self.model.last_decoder.feed_forward, PositionwiseFeedForwardDecoderSANM):
+            self.model.last_decoder.feed_forward = PositionwiseFeedForwardDecoderSANMExport(
+                self.model.last_decoder.feed_forward
+            )
+        self.last_decoder = self.model.last_decoder
+        self.bias_output = self.model.bias_output
+        self.dropout = self.model.dropout
+
+    def prepare_mask(self, mask):
+        """Prepare mask.
+        
+            Args:
+                mask: TODO.
+            """
+        mask_3d_btd = mask[:, :, None]
+        if len(mask.shape) == 2:
+            mask_4d_bhlt = 1 - mask[:, None, None, :]
+        elif len(mask.shape) == 3:
+            mask_4d_bhlt = 1 - mask[:, None, :]
+        mask_4d_bhlt = mask_4d_bhlt * -10000.0
+
+        return mask_3d_btd, mask_4d_bhlt
+
+    def forward(
+        self,
+        hs_pad: torch.Tensor,
+        hlens: torch.Tensor,
+        ys_in_pad: torch.Tensor,
+        ys_in_lens: torch.Tensor,
+        bias_embed: torch.Tensor,
+    ):
+
+        """Forward pass for training.
+        
+            Args:
+                hs_pad: TODO.
+                hlens: TODO.
+                ys_in_pad: TODO.
+                ys_in_lens: Lengths of ys_in.
+                bias_embed: TODO.
+            """
+        tgt = ys_in_pad
+        tgt_mask = self.make_pad_mask(ys_in_lens)
+        tgt_mask, _ = self.prepare_mask(tgt_mask)
+        # tgt_mask = myutils.sequence_mask(ys_in_lens, device=tgt.device)[:, :, None]
+
+        memory = hs_pad
+        memory_mask = self.make_pad_mask(hlens)
+        _, memory_mask = self.prepare_mask(memory_mask)
+        # memory_mask = myutils.sequence_mask(hlens, device=memory.device)[:, None, :]
+
+        x = tgt
+        x, tgt_mask, memory, memory_mask, _ = self.model.decoders(x, tgt_mask, memory, memory_mask)
+
+        _, _, x_self_attn, x_src_attn = self.last_decoder(x, tgt_mask, memory, memory_mask)
+
+        # contextual paraformer related
+        contextual_length = torch.Tensor([bias_embed.shape[1]]).int().repeat(hs_pad.shape[0])
+        # contextual_mask = myutils.sequence_mask(contextual_length, device=memory.device)[:, None, :]
+        contextual_mask = self.make_pad_mask(contextual_length)
+        contextual_mask, _ = self.prepare_mask(contextual_mask)
+        contextual_mask = contextual_mask.transpose(2, 1).unsqueeze(1)
+        cx, tgt_mask, _, _, _ = self.bias_decoder(
+            x_self_attn, tgt_mask, bias_embed, memory_mask=contextual_mask
+        )
+
+        if self.bias_output is not None:
+            x = torch.cat([x_src_attn, cx], dim=2)
+            x = self.bias_output(x.transpose(1, 2)).transpose(1, 2)  # 2D -> D
+            x = x_self_attn + self.dropout(x)
+
+        if self.model.decoders2 is not None:
+            x, tgt_mask, memory, memory_mask, _ = self.model.decoders2(
+                x, tgt_mask, memory, memory_mask
+            )
+        x, tgt_mask, memory, memory_mask, _ = self.model.decoders3(x, tgt_mask, memory, memory_mask)
+        x = self.after_norm(x)
+        x = self.output_layer(x)
+
+        return x, ys_in_lens

funasr/models/contextual_paraformer/export_meta.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import torch
+import types
+
+from funasr.register import tables
+from funasr.models.seaco_paraformer.export_meta import ContextualEmbedderExport
+
+
+class ContextualEmbedderExport2(ContextualEmbedderExport):
+    def __init__(self, model, **kwargs):
+        """Initialize ContextualEmbedderExport2.
+        
+            Args:
+                model: Model instance or model name.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__(model)
+        self.embedding = model.bias_embed
+        model.bias_encoder.batch_first = False
+        self.bias_encoder = model.bias_encoder
+    
+    def export_dummy_inputs(self):
+        """Export dummy inputs."""
+        hotword = torch.tensor(
+            [
+                [10, 11, 12, 13, 14, 10, 11, 12, 13, 14],
+                [100, 101, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [10, 11, 12, 13, 14, 10, 11, 12, 13, 14],
+                [100, 101, 0, 0, 0, 0, 0, 0, 0, 0],
+                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+            ],
+            dtype=torch.int32,
+        )
+        # hotword_length = torch.tensor([10, 2, 1], dtype=torch.int32)
+        return (hotword)
+
+
+def export_rebuild_model(model, **kwargs):
+    """Export rebuild model.
+    
+        Args:
+            model: Model instance or model name.
+            **kwargs: Additional keyword arguments.
+        """
+    is_onnx = kwargs.get("type", "onnx") == "onnx"
+
+    encoder_class = tables.encoder_classes.get(kwargs["encoder"] + "Export")
+    model.encoder = encoder_class(model.encoder, onnx=is_onnx)
+
+    predictor_class = tables.predictor_classes.get(kwargs["predictor"] + "Export")
+    model.predictor = predictor_class(model.predictor, onnx=is_onnx)
+
+    # little difference with bias encoder with seaco paraformer
+    embedder_class = ContextualEmbedderExport2
+    embedder_model = embedder_class(model, onnx=is_onnx)
+
+    if kwargs["decoder"] == "ParaformerSANMDecoder":
+        kwargs["decoder"] = "ParaformerSANMDecoderOnline"
+    decoder_class = tables.decoder_classes.get(kwargs["decoder"] + "Export")
+    model.decoder = decoder_class(model.decoder, onnx=is_onnx)
+
+    from funasr.utils.torch_function import sequence_mask
+
+    model.make_pad_mask = sequence_mask(kwargs["max_seq_len"], flip=False)
+    model.feats_dim = 560
+
+    import copy
+
+    backbone_model = copy.copy(model)
+
+    # backbone
+    backbone_model.forward = types.MethodType(export_backbone_forward, backbone_model)
+    backbone_model.export_dummy_inputs = types.MethodType(
+        export_backbone_dummy_inputs, backbone_model
+    )
+    backbone_model.export_input_names = types.MethodType(
+        export_backbone_input_names, backbone_model
+    )
+    backbone_model.export_output_names = types.MethodType(
+        export_backbone_output_names, backbone_model
+    )
+    backbone_model.export_dynamic_axes = types.MethodType(
+        export_backbone_dynamic_axes, backbone_model
+    )
+    
+    embedder_model.export_name = "model_eb"
+    backbone_model.export_name = "model"
+
+    return backbone_model, embedder_model
+
+
+def export_backbone_forward(
+    self,
+    speech: torch.Tensor,
+    speech_lengths: torch.Tensor,
+    bias_embed: torch.Tensor,
+):
+    """Export backbone forward.
+    
+        Args:
+            speech: Speech audio tensor, shape (batch, time).
+            speech_lengths: Length of each speech sample.
+            bias_embed: TODO.
+        """
+    batch = {"speech": speech, "speech_lengths": speech_lengths}
+
+    enc, enc_len = self.encoder(**batch)
+    mask = self.make_pad_mask(enc_len)[:, None, :]
+    pre_acoustic_embeds, pre_token_length, _, _ = self.predictor(enc, mask)
+    pre_token_length = pre_token_length.floor().type(torch.int32)
+
+    decoder_out, _ = self.decoder(enc, enc_len, pre_acoustic_embeds, pre_token_length, bias_embed)
+    decoder_out = torch.log_softmax(decoder_out, dim=-1)
+
+    return decoder_out, pre_token_length
+
+
+def export_backbone_dummy_inputs(self):
+    """Export backbone dummy inputs."""
+    speech = torch.randn(2, 30, self.feats_dim)
+    speech_lengths = torch.tensor([6, 30], dtype=torch.int32)
+    bias_embed = torch.randn(2, 1, 512)
+    return (speech, speech_lengths, bias_embed)
+
+
+def export_backbone_input_names(self):
+    """Export backbone input names."""
+    return ["speech", "speech_lengths", "bias_embed"]
+
+
+def export_backbone_output_names(self):
+    """Export backbone output names."""
+    return ["logits", "token_num"]
+
+
+def export_backbone_dynamic_axes(self):
+    """Export backbone dynamic axes."""
+    return {
+        "speech": {0: "batch_size", 1: "feats_length"},
+        "speech_lengths": {
+            0: "batch_size",
+        },
+        "bias_embed": {0: "batch_size", 1: "num_hotwords"},
+        "logits": {0: "batch_size", 1: "logits_length"},
+    }
+
+
+def export_backbone_name(self):
+    """Export backbone name."""
+    return "model.onnx"

funasr/models/contextual_paraformer/model.py

@@ -0,0 +1,673 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+
+import os
+import re
+import time
+import torch
+import codecs
+import logging
+import tempfile
+import requests
+import numpy as np
+from typing import Dict, Tuple
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.utils import postprocess_utils
+from funasr.metrics.compute_acc import th_accuracy
+from funasr.models.paraformer.model import Paraformer
+from funasr.utils.datadir_writer import DatadirWriter
+from funasr.models.paraformer.search import Hypothesis
+from funasr.train_utils.device_funcs import force_gatherable
+from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
+from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
+from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
+
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+@tables.register("model_classes", "ContextualParaformer")
+class ContextualParaformer(Paraformer):
+    """ContextualParaformer: Paraformer with hotword/context biasing.
+
+    Extends Paraformer with a context encoder that incorporates user-defined
+    hotwords/keywords to boost recognition of domain-specific terms.
+
+    Usage: Pass hotwords via generate(hotword='term1 term2').
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        """Initialize ContextualParaformer.
+        
+            Args:
+                *args: Variable positional arguments.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__(*args, **kwargs)
+
+        self.target_buffer_length = kwargs.get("target_buffer_length", -1)
+        inner_dim = kwargs.get("inner_dim", 256)
+        bias_encoder_type = kwargs.get("bias_encoder_type", "lstm")
+        use_decoder_embedding = kwargs.get("use_decoder_embedding", False)
+        crit_attn_weight = kwargs.get("crit_attn_weight", 0.0)
+        crit_attn_smooth = kwargs.get("crit_attn_smooth", 0.0)
+        bias_encoder_dropout_rate = kwargs.get("bias_encoder_dropout_rate", 0.0)
+
+        if bias_encoder_type == "lstm":
+            self.bias_encoder = torch.nn.LSTM(
+                inner_dim, inner_dim, 1, batch_first=True, dropout=bias_encoder_dropout_rate
+            )
+            self.bias_embed = torch.nn.Embedding(self.vocab_size, inner_dim)
+        elif bias_encoder_type == "mean":
+            self.bias_embed = torch.nn.Embedding(self.vocab_size, inner_dim)
+        else:
+            logging.error("Unsupport bias encoder type: {}".format(bias_encoder_type))
+
+        if self.target_buffer_length > 0:
+            self.hotword_buffer = None
+            self.length_record = []
+            self.current_buffer_length = 0
+        self.use_decoder_embedding = use_decoder_embedding
+        self.crit_attn_weight = crit_attn_weight
+        if self.crit_attn_weight > 0:
+            self.attn_loss = torch.nn.L1Loss()
+        self.crit_attn_smooth = crit_attn_smooth
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
+        """Frontend + Encoder + Decoder + Calc loss
+
+        Args:
+                speech: (Batch, Length, ...)
+                speech_lengths: (Batch, )
+                text: (Batch, Length)
+                text_lengths: (Batch,)
+        """
+        text_lengths = text_lengths.squeeze()
+        speech_lengths = speech_lengths.squeeze()
+
+        batch_size = speech.shape[0]
+
+        hotword_pad = kwargs.get("hotword_pad")
+        hotword_lengths = kwargs.get("hotword_lengths")
+        # dha_pad = kwargs.get("dha_pad")
+
+        # 1. Encoder
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+
+        loss_ctc, cer_ctc = None, None
+
+        stats = dict()
+
+        # 1. CTC branch
+        if self.ctc_weight != 0.0:
+            loss_ctc, cer_ctc = self._calc_ctc_loss(
+                encoder_out, encoder_out_lens, text, text_lengths
+            )
+
+            # Collect CTC branch stats
+            stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
+            stats["cer_ctc"] = cer_ctc
+
+        # 2b. Attention decoder branch
+        loss_att, acc_att, cer_att, wer_att, loss_pre, loss_ideal = self._calc_att_clas_loss(
+            encoder_out, encoder_out_lens, text, text_lengths, hotword_pad, hotword_lengths
+        )
+
+        # 3. CTC-Att loss definition
+        if self.ctc_weight == 0.0:
+            loss = loss_att + loss_pre * self.predictor_weight
+        else:
+            loss = (
+                self.ctc_weight * loss_ctc
+                + (1 - self.ctc_weight) * loss_att
+                + loss_pre * self.predictor_weight
+            )
+
+        if loss_ideal is not None:
+            loss = loss + loss_ideal * self.crit_attn_weight
+            stats["loss_ideal"] = loss_ideal.detach().cpu()
+
+        # Collect Attn branch stats
+        stats["loss_att"] = loss_att.detach() if loss_att is not None else None
+        stats["acc"] = acc_att
+        stats["cer"] = cer_att
+        stats["wer"] = wer_att
+        stats["loss_pre"] = loss_pre.detach().cpu() if loss_pre is not None else None
+
+        stats["loss"] = torch.clone(loss.detach())
+        # force_gatherable: to-device and to-tensor if scalar for DataParallel
+        if self.length_normalized_loss:
+            batch_size = int((text_lengths + self.predictor_bias).sum())
+
+        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
+        return loss, stats, weight
+
+    def _calc_att_clas_loss(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_out_lens: torch.Tensor,
+        ys_pad: torch.Tensor,
+        ys_pad_lens: torch.Tensor,
+        hotword_pad: torch.Tensor,
+        hotword_lengths: torch.Tensor,
+    ):
+        """Internal: calc att clas loss.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+                hotword_pad: TODO.
+                hotword_lengths: Lengths of hotword.
+            """
+        encoder_out_mask = (
+            ~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]
+        ).to(encoder_out.device)
+
+        if self.predictor_bias == 1:
+            _, ys_pad = add_sos_eos(ys_pad, self.sos, self.eos, self.ignore_id)
+            ys_pad_lens = ys_pad_lens + self.predictor_bias
+
+        pre_acoustic_embeds, pre_token_length, _, _ = self.predictor(
+            encoder_out, ys_pad, encoder_out_mask, ignore_id=self.ignore_id
+        )
+        # -1. bias encoder
+        if self.use_decoder_embedding:
+            hw_embed = self.decoder.embed(hotword_pad)
+        else:
+            hw_embed = self.bias_embed(hotword_pad)
+
+        hw_embed, (_, _) = self.bias_encoder(hw_embed)
+        _ind = np.arange(0, hotword_pad.shape[0]).tolist()
+        selected = hw_embed[_ind, [i - 1 for i in hotword_lengths.detach().cpu().tolist()]]
+        contextual_info = selected.squeeze(0).repeat(ys_pad.shape[0], 1, 1).to(ys_pad.device)
+
+        # 0. sampler
+        decoder_out_1st = None
+        if self.sampling_ratio > 0.0:
+
+            sematic_embeds, decoder_out_1st = self.sampler(
+                encoder_out,
+                encoder_out_lens,
+                ys_pad,
+                ys_pad_lens,
+                pre_acoustic_embeds,
+                contextual_info,
+            )
+        else:
+            sematic_embeds = pre_acoustic_embeds
+
+        # 1. Forward decoder
+        decoder_outs = self.decoder(
+            encoder_out,
+            encoder_out_lens,
+            sematic_embeds,
+            ys_pad_lens,
+            contextual_info=contextual_info,
+        )
+        decoder_out, _ = decoder_outs[0], decoder_outs[1]
+        """
+        if self.crit_attn_weight > 0 and attn.shape[-1] > 1:
+            ideal_attn = ideal_attn + self.crit_attn_smooth / (self.crit_attn_smooth + 1.0)
+            attn_non_blank = attn[:,:,:,:-1]
+            ideal_attn_non_blank = ideal_attn[:,:,:-1]
+            loss_ideal = self.attn_loss(attn_non_blank.max(1)[0], ideal_attn_non_blank.to(attn.device))
+        else:
+            loss_ideal = None
+        """
+        loss_ideal = None
+
+        if decoder_out_1st is None:
+            decoder_out_1st = decoder_out
+        # 2. Compute attention loss
+        loss_att = self.criterion_att(decoder_out, ys_pad)
+        acc_att = th_accuracy(
+            decoder_out_1st.view(-1, self.vocab_size),
+            ys_pad,
+            ignore_label=self.ignore_id,
+        )
+        loss_pre = self.criterion_pre(ys_pad_lens.type_as(pre_token_length), pre_token_length)
+
+        # Compute cer/wer using attention-decoder
+        if self.training or self.error_calculator is None:
+            cer_att, wer_att = None, None
+        else:
+            ys_hat = decoder_out_1st.argmax(dim=-1)
+            cer_att, wer_att = self.error_calculator(ys_hat.cpu(), ys_pad.cpu())
+
+        return loss_att, acc_att, cer_att, wer_att, loss_pre, loss_ideal
+
+    def sampler(
+        self,
+        encoder_out,
+        encoder_out_lens,
+        ys_pad,
+        ys_pad_lens,
+        pre_acoustic_embeds,
+        contextual_info,
+    ):
+        """Sampler.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                ys_pad: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+                pre_acoustic_embeds: TODO.
+                contextual_info: TODO.
+            """
+        tgt_mask = (~make_pad_mask(ys_pad_lens, maxlen=ys_pad_lens.max())[:, :, None]).to(
+            ys_pad.device
+        )
+        ys_pad = ys_pad * tgt_mask[:, :, 0]
+        if self.share_embedding:
+            ys_pad_embed = self.decoder.output_layer.weight[ys_pad]
+        else:
+            ys_pad_embed = self.decoder.embed(ys_pad)
+        with torch.no_grad():
+            decoder_outs = self.decoder(
+                encoder_out,
+                encoder_out_lens,
+                pre_acoustic_embeds,
+                ys_pad_lens,
+                contextual_info=contextual_info,
+            )
+            decoder_out, _ = decoder_outs[0], decoder_outs[1]
+            pred_tokens = decoder_out.argmax(-1)
+            nonpad_positions = ys_pad.ne(self.ignore_id)
+            seq_lens = (nonpad_positions).sum(1)
+            same_num = ((pred_tokens == ys_pad) & nonpad_positions).sum(1)
+            input_mask = torch.ones_like(nonpad_positions)
+            bsz, seq_len = ys_pad.size()
+            for li in range(bsz):
+                target_num = (
+                    ((seq_lens[li] - same_num[li].sum()).float()) * self.sampling_ratio
+                ).long()
+                if target_num > 0:
+                    input_mask[li].scatter_(
+                        dim=0,
+                        index=torch.randperm(seq_lens[li])[:target_num].to(
+                            pre_acoustic_embeds.device
+                        ),
+                        value=0,
+                    )
+            input_mask = input_mask.eq(1)
+            input_mask = input_mask.masked_fill(~nonpad_positions, False)
+            input_mask_expand_dim = input_mask.unsqueeze(2).to(pre_acoustic_embeds.device)
+
+        sematic_embeds = pre_acoustic_embeds.masked_fill(
+            ~input_mask_expand_dim, 0
+        ) + ys_pad_embed.masked_fill(input_mask_expand_dim, 0)
+        return sematic_embeds * tgt_mask, decoder_out * tgt_mask
+
+    def cal_decoder_with_predictor(
+        self,
+        encoder_out,
+        encoder_out_lens,
+        sematic_embeds,
+        ys_pad_lens,
+        hw_list=None,
+        clas_scale=1.0,
+    ):
+        """Cal decoder with predictor.
+        
+            Args:
+                encoder_out: Encoder output tensor.
+                encoder_out_lens: Encoder output lengths.
+                sematic_embeds: TODO.
+                ys_pad_lens: Lengths of ys_pad.
+                hw_list: TODO.
+                clas_scale: TODO.
+            """
+        if hw_list is None:
+            hw_list = [torch.Tensor([1]).long().to(encoder_out.device)]  # empty hotword list
+            hw_list_pad = pad_list(hw_list, 0)
+            if self.use_decoder_embedding:
+                hw_embed = self.decoder.embed(hw_list_pad)
+            else:
+                hw_embed = self.bias_embed(hw_list_pad)
+            hw_embed, (h_n, _) = self.bias_encoder(hw_embed)
+            hw_embed = h_n.repeat(encoder_out.shape[0], 1, 1)
+        else:
+            hw_lengths = [len(i) for i in hw_list]
+            hw_list_pad = pad_list([torch.Tensor(i).long() for i in hw_list], 0).to(
+                encoder_out.device
+            )
+            if self.use_decoder_embedding:
+                hw_embed = self.decoder.embed(hw_list_pad)
+            else:
+                hw_embed = self.bias_embed(hw_list_pad)
+            hw_embed = torch.nn.utils.rnn.pack_padded_sequence(
+                hw_embed, hw_lengths, batch_first=True, enforce_sorted=False
+            )
+            _, (h_n, _) = self.bias_encoder(hw_embed)
+            hw_embed = h_n.repeat(encoder_out.shape[0], 1, 1)
+
+        decoder_outs = self.decoder(
+            encoder_out,
+            encoder_out_lens,
+            sematic_embeds,
+            ys_pad_lens,
+            contextual_info=hw_embed,
+            clas_scale=clas_scale,
+        )
+
+        decoder_out = decoder_outs[0]
+        decoder_out = torch.log_softmax(decoder_out, dim=-1)
+        return decoder_out, ys_pad_lens
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        # init beamsearch
+
+        """Run inference on input data.
+        
+            Args:
+                data_in: Input data (audio samples, file paths, or text).
+                data_lengths: Lengths of each input sample in the batch.
+                key: Sample identifiers.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                frontend: Audio frontend for feature extraction.
+                **kwargs: Additional keyword arguments.
+            """
+        is_use_ctc = kwargs.get("decoding_ctc_weight", 0.0) > 0.00001 and self.ctc != None
+        is_use_lm = (
+            kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None
+        )
+        if self.beam_search is None and (is_use_lm or is_use_ctc):
+            logging.info("enable beam_search")
+            self.init_beam_search(**kwargs)
+            self.nbest = kwargs.get("nbest", 1)
+
+        meta_data = {}
+
+        # extract fbank feats
+        time1 = time.perf_counter()
+
+        audio_sample_list = load_audio_text_image_video(
+            data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000)
+        )
+
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+
+        speech, speech_lengths = extract_fbank(
+            audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
+        )
+        time3 = time.perf_counter()
+        meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+        meta_data["batch_data_time"] = (
+            speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+        )
+
+        speech = speech.to(device=kwargs["device"])
+        speech_lengths = speech_lengths.to(device=kwargs["device"])
+
+        # hotword
+        self.hotword_list = self.generate_hotwords_list(
+            kwargs.get("hotword", None), tokenizer=tokenizer, frontend=frontend
+        )
+
+        # Encoder
+        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)
+        if isinstance(encoder_out, tuple):
+            encoder_out = encoder_out[0]
+
+        # predictor
+        predictor_outs = self.calc_predictor(encoder_out, encoder_out_lens)
+        pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = (
+            predictor_outs[0],
+            predictor_outs[1],
+            predictor_outs[2],
+            predictor_outs[3],
+        )
+        pre_token_length = pre_token_length.round().long()
+        if torch.max(pre_token_length) < 1:
+            return []
+
+        decoder_outs = self.cal_decoder_with_predictor(
+            encoder_out,
+            encoder_out_lens,
+            pre_acoustic_embeds,
+            pre_token_length,
+            hw_list=self.hotword_list,
+            clas_scale=kwargs.get("clas_scale", 1.0),
+        )
+        decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1]
+
+        results = []
+        b, n, d = decoder_out.size()
+        for i in range(b):
+            x = encoder_out[i, : encoder_out_lens[i], :]
+            am_scores = decoder_out[i, : pre_token_length[i], :]
+            if self.beam_search is not None:
+                nbest_hyps = self.beam_search(
+                    x=x,
+                    am_scores=am_scores,
+                    maxlenratio=kwargs.get("maxlenratio", 0.0),
+                    minlenratio=kwargs.get("minlenratio", 0.0),
+                )
+
+                nbest_hyps = nbest_hyps[: self.nbest]
+            else:
+
+                yseq = am_scores.argmax(dim=-1)
+                score = am_scores.max(dim=-1)[0]
+                score = torch.sum(score, dim=-1)
+                # pad with mask tokens to ensure compatibility with sos/eos tokens
+                yseq = torch.tensor([self.sos] + yseq.tolist() + [self.eos], device=yseq.device)
+                nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
+            for nbest_idx, hyp in enumerate(nbest_hyps):
+                ibest_writer = None
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
+
+                # remove sos/eos and get results
+                last_pos = -1
+                if isinstance(hyp.yseq, list):
+                    token_int = hyp.yseq[1:last_pos]
+                else:
+                    token_int = hyp.yseq[1:last_pos].tolist()
+
+                # remove blank symbol id, which is assumed to be 0
+                token_int = list(
+                    filter(
+                        lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int
+                    )
+                )
+
+                if tokenizer is not None:
+                    # Change integer-ids to tokens
+                    token = tokenizer.ids2tokens(token_int)
+                    text = tokenizer.tokens2text(token)
+
+                    text_postprocessed, _ = postprocess_utils.sentence_postprocess(token)
+                    result_i = {"key": key[i], "text": text_postprocessed}
+
+                    if ibest_writer is not None:
+                        ibest_writer["token"][key[i]] = " ".join(token)
+                        ibest_writer["text"][key[i]] = text
+                        ibest_writer["text_postprocessed"][key[i]] = text_postprocessed
+                else:
+                    result_i = {"key": key[i], "token_int": token_int}
+                results.append(result_i)
+
+        return results, meta_data
+
+    def generate_hotwords_list(self, hotword_list_or_file, tokenizer=None, frontend=None):
+        """Generate hotwords list.
+        
+            Args:
+                hotword_list_or_file: TODO.
+                tokenizer: Tokenizer instance for text encoding/decoding.
+                frontend: Audio frontend for feature extraction.
+            """
+        def load_seg_dict(seg_dict_file):
+            """Load seg dict.
+            
+                Args:
+                    seg_dict_file: TODO.
+                """
+            seg_dict = {}
+            assert isinstance(seg_dict_file, str)
+            with open(seg_dict_file, "r", encoding="utf8") as f:
+                lines = f.readlines()
+                for line in lines:
+                    s = line.strip().split()
+                    key = s[0]
+                    value = s[1:]
+                    seg_dict[key] = " ".join(value)
+            return seg_dict
+
+        def seg_tokenize(txt, seg_dict):
+            """Seg tokenize.
+            
+                Args:
+                    txt: TODO.
+                    seg_dict: TODO.
+                """
+            pattern = re.compile(r"^[\u4E00-\u9FA50-9]+$")
+            out_txt = ""
+            for word in txt:
+                word = word.lower()
+                if word in seg_dict:
+                    out_txt += seg_dict[word] + " "
+                else:
+                    if pattern.match(word):
+                        for char in word:
+                            if char in seg_dict:
+                                out_txt += seg_dict[char] + " "
+                            else:
+                                out_txt += "<unk>" + " "
+                    else:
+                        out_txt += "<unk>" + " "
+            return out_txt.strip().split()
+
+        seg_dict = None
+        if frontend.cmvn_file is not None:
+            model_dir = os.path.dirname(frontend.cmvn_file)
+            seg_dict_file = os.path.join(model_dir, "seg_dict")
+            if os.path.exists(seg_dict_file):
+                seg_dict = load_seg_dict(seg_dict_file)
+            else:
+                seg_dict = None
+        # for None
+        if hotword_list_or_file is None:
+            hotword_list = None
+        # for local txt inputs
+        elif os.path.exists(hotword_list_or_file) and hotword_list_or_file.endswith(".txt"):
+            logging.info("Attempting to parse hotwords from local txt...")
+            hotword_list = []
+            hotword_str_list = []
+            with codecs.open(hotword_list_or_file, "r") as fin:
+                for line in fin.readlines():
+                    hw = line.strip()
+                    hw_list = hw.split()
+                    if seg_dict is not None:
+                        hw_list = seg_tokenize(hw_list, seg_dict)
+                    hotword_str_list.append(hw)
+                    hotword_list.append(tokenizer.tokens2ids(hw_list))
+                hotword_list.append([self.sos])
+                hotword_str_list.append("<s>")
+            logging.info(
+                "Initialized hotword list from file: {}, hotword list: {}.".format(
+                    hotword_list_or_file, hotword_str_list
+                )
+            )
+        # for url, download and generate txt
+        elif hotword_list_or_file.startswith("http"):
+            logging.info("Attempting to parse hotwords from url...")
+            work_dir = tempfile.TemporaryDirectory().name
+            if not os.path.exists(work_dir):
+                os.makedirs(work_dir)
+            text_file_path = os.path.join(work_dir, os.path.basename(hotword_list_or_file))
+            local_file = requests.get(hotword_list_or_file)
+            open(text_file_path, "wb").write(local_file.content)
+            hotword_list_or_file = text_file_path
+            hotword_list = []
+            hotword_str_list = []
+            with codecs.open(hotword_list_or_file, "r") as fin:
+                for line in fin.readlines():
+                    hw = line.strip()
+                    hw_list = hw.split()
+                    if seg_dict is not None:
+                        hw_list = seg_tokenize(hw_list, seg_dict)
+                    hotword_str_list.append(hw)
+                    hotword_list.append(tokenizer.tokens2ids(hw_list))
+                hotword_list.append([self.sos])
+                hotword_str_list.append("<s>")
+            logging.info(
+                "Initialized hotword list from file: {}, hotword list: {}.".format(
+                    hotword_list_or_file, hotword_str_list
+                )
+            )
+        # for text str input
+        elif not hotword_list_or_file.endswith(".txt"):
+            logging.info("Attempting to parse hotwords as str...")
+            hotword_list = []
+            hotword_str_list = []
+            for hw in hotword_list_or_file.strip().split():
+                hotword_str_list.append(hw)
+                hw_list = hw.strip().split()
+                if seg_dict is not None:
+                    hw_list = seg_tokenize(hw_list, seg_dict)
+                hotword_list.append(tokenizer.tokens2ids(hw_list))
+            hotword_list.append([self.sos])
+            hotword_str_list.append("<s>")
+            logging.info("Hotword list: {}.".format(hotword_str_list))
+        else:
+            hotword_list = None
+        return hotword_list
+
+    def export(
+        self,
+        **kwargs,
+    ):
+        """Export.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        if "max_seq_len" not in kwargs:
+            kwargs["max_seq_len"] = 512
+        from .export_meta import export_rebuild_model
+
+        models = export_rebuild_model(model=self, **kwargs)
+        return models

funasr/models/contextual_paraformer/template.yaml

@@ -0,0 +1,129 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: ContextualParaformer
+model_conf:
+    ctc_weight: 0.0
+    lsm_weight: 0.1
+    length_normalized_loss: true
+    predictor_weight: 1.0
+    predictor_bias: 1
+    sampling_ratio: 0.75
+    inner_dim: 512
+
+# encoder
+encoder: SANMEncoder
+encoder_conf:
+    output_size: 512
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 50
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: pe
+    pos_enc_class: SinusoidalPositionEncoder
+    normalize_before: true
+    kernel_size: 11
+    sanm_shfit: 0
+    selfattention_layer_type: sanm
+
+
+# decoder
+decoder: ContextualParaformerDecoder
+decoder_conf:
+    attention_heads: 4
+    linear_units: 2048
+    num_blocks: 16
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    att_layer_num: 16
+    kernel_size: 11
+    sanm_shfit: 0
+
+predictor: CifPredictorV2
+predictor_conf:
+    idim: 512
+    threshold: 1.0
+    l_order: 1
+    r_order: 1
+    tail_threshold: 0.45
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000
+    window: hamming
+    n_mels: 80
+    frame_length: 25
+    frame_shift: 10
+    lfr_m: 7
+    lfr_n: 6
+
+specaug: SpecAugLFR
+specaug_conf:
+    apply_time_warp: false
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    lfr_rate: 6
+    num_freq_mask: 1
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 12
+    num_time_mask: 1
+
+train_conf:
+  accum_grad: 1
+  grad_clip: 5
+  max_epoch: 150
+  val_scheduler_criterion:
+      - valid
+      - acc
+  best_model_criterion:
+  -   - valid
+      - acc
+      - max
+  keep_nbest_models: 10
+  log_interval: 50
+
+optim: adam
+optim_conf:
+   lr: 0.0005
+scheduler: warmuplr
+scheduler_conf:
+   warmup_steps: 30000
+
+dataset: AudioDataset
+dataset_conf:
+    index_ds: IndexDSJsonl
+    batch_sampler: BatchSampler
+    batch_type: example # example or length
+    batch_size: 1 # if batch_type is example, batch_size is the numbers of samples; if length, batch_size is source_token_len+target_token_len;
+    max_token_length: 2048 # filter samples if source_token_len+target_token_len > max_token_length,
+    buffer_size: 500
+    shuffle: True
+    num_workers: 0
+
+tokenizer: CharTokenizer
+tokenizer_conf:
+  unk_symbol: <unk>
+  split_with_space: true
+
+ctc_conf:
+    dropout_rate: 0.0
+    ctc_type: builtin
+    reduce: true
+    ignore_nan_grad: true
+normalize: null