Upload 16 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tts-2025-04-27@197cd135f2a2451b9cab9cf2add1c1ab.wav filter=lfs diff=lfs merge=lfs -text

conformer/conformer/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .model import Conformer

conformer/conformer/activation.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch.nn as nn
+from torch import Tensor
+
+
+class Swish(nn.Module):
+    """
+    Swish is a smooth, non-monotonic function that consistently matches or outperforms ReLU on deep networks applied
+    to a variety of challenging domains such as Image classification and Machine translation.
+    """
+    def __init__(self):
+        super(Swish, self).__init__()
+    
+    def forward(self, inputs: Tensor) -> Tensor:
+        return inputs * inputs.sigmoid()
+
+
+class GLU(nn.Module):
+    """
+    The gating mechanism is called Gated Linear Units (GLU), which was first introduced for natural language processing
+    in the paper “Language Modeling with Gated Convolutional Networks”
+    """
+    def __init__(self, dim: int) -> None:
+        super(GLU, self).__init__()
+        self.dim = dim
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        outputs, gate = inputs.chunk(2, dim=self.dim)
+        return outputs * gate.sigmoid()

conformer/conformer/attention.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from typing import Optional
+
+from .embedding import PositionalEncoding
+from .modules import Linear
+
+
+class RelativeMultiHeadAttention(nn.Module):
+    """
+    Multi-head attention with relative positional encoding.
+    This concept was proposed in the "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context"
+
+    Args:
+        d_model (int): The dimension of model
+        num_heads (int): The number of attention heads.
+        dropout_p (float): probability of dropout
+
+    Inputs: query, key, value, pos_embedding, mask
+        - **query** (batch, time, dim): Tensor containing query vector
+        - **key** (batch, time, dim): Tensor containing key vector
+        - **value** (batch, time, dim): Tensor containing value vector
+        - **pos_embedding** (batch, time, dim): Positional embedding tensor
+        - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+
+    Returns:
+        - **outputs**: Tensor produces by relative multi head attention module.
+    """
+    def __init__(
+            self,
+            d_model: int = 512,
+            num_heads: int = 16,
+            dropout_p: float = 0.1,
+    ):
+        super(RelativeMultiHeadAttention, self).__init__()
+        assert d_model % num_heads == 0, "d_model % num_heads should be zero."
+        self.d_model = d_model
+        self.d_head = int(d_model / num_heads)
+        self.num_heads = num_heads
+        self.sqrt_dim = math.sqrt(d_model)
+
+        self.query_proj = Linear(d_model, d_model)
+        self.key_proj = Linear(d_model, d_model)
+        self.value_proj = Linear(d_model, d_model)
+        self.pos_proj = Linear(d_model, d_model, bias=False)
+
+        self.dropout = nn.Dropout(p=dropout_p)
+        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        torch.nn.init.xavier_uniform_(self.u_bias)
+        torch.nn.init.xavier_uniform_(self.v_bias)
+
+        self.out_proj = Linear(d_model, d_model)
+
+    def forward(
+            self,
+            query: Tensor,
+            key: Tensor,
+            value: Tensor,
+            pos_embedding: Tensor,
+            mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        batch_size = value.size(0)
+
+        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)
+        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)
+        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)
+        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)
+
+        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))
+        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))
+        pos_score = self._relative_shift(pos_score)
+
+        score = (content_score + pos_score) / self.sqrt_dim
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+            score.masked_fill_(mask, -1e9)
+
+        attn = F.softmax(score, -1)
+        attn = self.dropout(attn)
+
+        context = torch.matmul(attn, value).transpose(1, 2)
+        context = context.contiguous().view(batch_size, -1, self.d_model)
+
+        return self.out_proj(context)
+
+    def _relative_shift(self, pos_score: Tensor) -> Tensor:
+        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()
+        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)
+        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)
+
+        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)
+        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)
+
+        return pos_score
+
+
+class MultiHeadedSelfAttentionModule(nn.Module):
+    """
+    Conformer employ multi-headed self-attention (MHSA) while integrating an important technique from Transformer-XL,
+    the relative sinusoidal positional encoding scheme. The relative positional encoding allows the self-attention
+    module to generalize better on different input length and the resulting encoder is more robust to the variance of
+    the utterance length. Conformer use prenorm residual units with dropout which helps training
+    and regularizing deeper models.
+
+    Args:
+        d_model (int): The dimension of model
+        num_heads (int): The number of attention heads.
+        dropout_p (float): probability of dropout
+
+    Inputs: inputs, mask
+        - **inputs** (batch, time, dim): Tensor containing input vector
+        - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+
+    Returns:
+        - **outputs** (batch, time, dim): Tensor produces by relative multi headed self attention module.
+    """
+    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.1):
+        super(MultiHeadedSelfAttentionModule, self).__init__()
+        self.positional_encoding = PositionalEncoding(d_model)
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)
+        self.dropout = nn.Dropout(p=dropout_p)
+
+    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):
+        batch_size, seq_length, _ = inputs.size()
+        pos_embedding = self.positional_encoding(seq_length)
+        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)
+
+        inputs = self.layer_norm(inputs)
+        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)
+
+        return self.dropout(outputs)

conformer/conformer/convolution.py

@@ -0,0 +1,186 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Tuple
+
+from .activation import Swish, GLU
+from .modules import Transpose
+
+
+class DepthwiseConv1d(nn.Module):
+    """
+    When groups == in_channels and out_channels == K * in_channels, where K is a positive integer,
+    this operation is termed in literature as depthwise convolution.
+
+    Args:
+        in_channels (int): Number of channels in the input
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the convolving kernel
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        bias (bool, optional): If True, adds a learnable bias to the output. Default: True
+
+    Inputs: inputs
+        - **inputs** (batch, in_channels, time): Tensor containing input vector
+
+    Returns: outputs
+        - **outputs** (batch, out_channels, time): Tensor produces by depthwise 1-D convolution.
+    """
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: int,
+            stride: int = 1,
+            padding: int = 0,
+            bias: bool = False,
+    ) -> None:
+        super(DepthwiseConv1d, self).__init__()
+        assert out_channels % in_channels == 0, "out_channels should be constant multiple of in_channels"
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            groups=in_channels,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.conv(inputs)
+
+
+class PointwiseConv1d(nn.Module):
+    """
+    When kernel size == 1 conv1d, this operation is termed in literature as pointwise convolution.
+    This operation often used to match dimensions.
+
+    Args:
+        in_channels (int): Number of channels in the input
+        out_channels (int): Number of channels produced by the convolution
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
+        bias (bool, optional): If True, adds a learnable bias to the output. Default: True
+
+    Inputs: inputs
+        - **inputs** (batch, in_channels, time): Tensor containing input vector
+
+    Returns: outputs
+        - **outputs** (batch, out_channels, time): Tensor produces by pointwise 1-D convolution.
+    """
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            stride: int = 1,
+            padding: int = 0,
+            bias: bool = True,
+    ) -> None:
+        super(PointwiseConv1d, self).__init__()
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.conv(inputs)
+
+
+class ConformerConvModule(nn.Module):
+    """
+    Conformer convolution module starts with a pointwise convolution and a gated linear unit (GLU).
+    This is followed by a single 1-D depthwise convolution layer. Batchnorm is  deployed just after the convolution
+    to aid training deep models.
+
+    Args:
+        in_channels (int): Number of channels in the input
+        kernel_size (int or tuple, optional): Size of the convolving kernel Default: 31
+        dropout_p (float, optional): probability of dropout
+
+    Inputs: inputs
+        inputs (batch, time, dim): Tensor contains input sequences
+
+    Outputs: outputs
+        outputs (batch, time, dim): Tensor produces by conformer convolution module.
+    """
+    def __init__(
+            self,
+            in_channels: int,
+            kernel_size: int = 31,
+            expansion_factor: int = 2,
+            dropout_p: float = 0.1,
+    ) -> None:
+        super(ConformerConvModule, self).__init__()
+        assert (kernel_size - 1) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding"
+        assert expansion_factor == 2, "Currently, Only Supports expansion_factor 2"
+
+        self.sequential = nn.Sequential(
+            nn.LayerNorm(in_channels),
+            Transpose(shape=(1, 2)),
+            PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True),
+            GLU(dim=1),
+            DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2),
+            nn.BatchNorm1d(in_channels),
+            Swish(),
+            PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True),
+            nn.Dropout(p=dropout_p),
+        )
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.sequential(inputs).transpose(1, 2)
+
+
+class Conv2dSubampling(nn.Module):
+    """
+    Convolutional 2D subsampling (to 1/4 length)
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+
+    Inputs: inputs
+        - **inputs** (batch, time, dim): Tensor containing sequence of inputs
+
+    Returns: outputs, output_lengths
+        - **outputs** (batch, time, dim): Tensor produced by the convolution
+        - **output_lengths** (batch): list of sequence output lengths
+    """
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super(Conv2dSubampling, self).__init__()
+        self.sequential = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2),
+            nn.ReLU(),
+        )
+
+    def forward(self, inputs: Tensor, input_lengths: Tensor) -> Tuple[Tensor, Tensor]:
+        outputs = self.sequential(inputs.unsqueeze(1))
+        batch_size, channels, subsampled_lengths, sumsampled_dim = outputs.size()
+
+        outputs = outputs.permute(0, 2, 1, 3)
+        outputs = outputs.contiguous().view(batch_size, subsampled_lengths, channels * sumsampled_dim)
+
+        output_lengths = input_lengths >> 2
+        output_lengths -= 1
+
+        return outputs, output_lengths

conformer/conformer/embedding.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class PositionalEncoding(nn.Module):
+    """
+    Positional Encoding proposed in "Attention Is All You Need".
+    Since transformer contains no recurrence and no convolution, in order for the model to make
+    use of the order of the sequence, we must add some positional information.
+
+    "Attention Is All You Need" use sine and cosine functions of different frequencies:
+        PE_(pos, 2i)    =  sin(pos / power(10000, 2i / d_model))
+        PE_(pos, 2i+1)  =  cos(pos / power(10000, 2i / d_model))
+    """
+    def __init__(self, d_model: int = 512, max_len: int = 10000) -> None:
+        super(PositionalEncoding, self).__init__()
+        pe = torch.zeros(max_len, d_model, requires_grad=False)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, length: int) -> Tensor:
+        return self.pe[:, :length]

conformer/conformer/encoder.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Tuple
+
+from .feed_forward import FeedForwardModule
+from .attention import MultiHeadedSelfAttentionModule
+from .convolution import (
+    ConformerConvModule,
+    Conv2dSubampling,
+)
+from .modules import (
+    ResidualConnectionModule,
+    Linear,
+)
+
+
+class ConformerBlock(nn.Module):
+    """
+    Conformer block contains two Feed Forward modules sandwiching the Multi-Headed Self-Attention module
+    and the Convolution module. This sandwich structure is inspired by Macaron-Net, which proposes replacing
+    the original feed-forward layer in the Transformer block into two half-step feed-forward layers,
+    one before the attention layer and one after.
+
+    Args:
+        encoder_dim (int, optional): Dimension of conformer encoder
+        num_attention_heads (int, optional): Number of attention heads
+        feed_forward_expansion_factor (int, optional): Expansion factor of feed forward module
+        conv_expansion_factor (int, optional): Expansion factor of conformer convolution module
+        feed_forward_dropout_p (float, optional): Probability of feed forward module dropout
+        attention_dropout_p (float, optional): Probability of attention module dropout
+        conv_dropout_p (float, optional): Probability of conformer convolution module dropout
+        conv_kernel_size (int or tuple, optional): Size of the convolving kernel
+        half_step_residual (bool): Flag indication whether to use half step residual or not
+
+    Inputs: inputs
+        - **inputs** (batch, time, dim): Tensor containing input vector
+
+    Returns: outputs
+        - **outputs** (batch, time, dim): Tensor produces by conformer block.
+    """
+    def __init__(
+            self,
+            encoder_dim: int = 512,
+            num_attention_heads: int = 8,
+            feed_forward_expansion_factor: int = 4,
+            conv_expansion_factor: int = 2,
+            feed_forward_dropout_p: float = 0.1,
+            attention_dropout_p: float = 0.1,
+            conv_dropout_p: float = 0.1,
+            conv_kernel_size: int = 31,
+            half_step_residual: bool = True,
+    ):
+        super(ConformerBlock, self).__init__()
+        if half_step_residual:
+            self.feed_forward_residual_factor = 0.5
+        else:
+            self.feed_forward_residual_factor = 1
+
+        self.sequential = nn.Sequential(
+            ResidualConnectionModule(
+                module=FeedForwardModule(
+                    encoder_dim=encoder_dim,
+                    expansion_factor=feed_forward_expansion_factor,
+                    dropout_p=feed_forward_dropout_p,
+                ),
+                module_factor=self.feed_forward_residual_factor,
+            ),
+            ResidualConnectionModule(
+                module=MultiHeadedSelfAttentionModule(
+                    d_model=encoder_dim,
+                    num_heads=num_attention_heads,
+                    dropout_p=attention_dropout_p,
+                ),
+            ),
+            ResidualConnectionModule(
+                module=ConformerConvModule(
+                    in_channels=encoder_dim,
+                    kernel_size=conv_kernel_size,
+                    expansion_factor=conv_expansion_factor,
+                    dropout_p=conv_dropout_p,
+                ),
+            ),
+            ResidualConnectionModule(
+                module=FeedForwardModule(
+                    encoder_dim=encoder_dim,
+                    expansion_factor=feed_forward_expansion_factor,
+                    dropout_p=feed_forward_dropout_p,
+                ),
+                module_factor=self.feed_forward_residual_factor,
+            ),
+            nn.LayerNorm(encoder_dim),
+        )
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.sequential(inputs)
+
+
+class ConformerEncoder(nn.Module):
+    """
+    Conformer encoder first processes the input with a convolution subsampling layer and then
+    with a number of conformer blocks.
+
+    Args:
+        input_dim (int, optional): Dimension of input vector
+        encoder_dim (int, optional): Dimension of conformer encoder
+        num_layers (int, optional): Number of conformer blocks
+        num_attention_heads (int, optional): Number of attention heads
+        feed_forward_expansion_factor (int, optional): Expansion factor of feed forward module
+        conv_expansion_factor (int, optional): Expansion factor of conformer convolution module
+        feed_forward_dropout_p (float, optional): Probability of feed forward module dropout
+        attention_dropout_p (float, optional): Probability of attention module dropout
+        conv_dropout_p (float, optional): Probability of conformer convolution module dropout
+        conv_kernel_size (int or tuple, optional): Size of the convolving kernel
+        half_step_residual (bool): Flag indication whether to use half step residual or not
+
+    Inputs: inputs, input_lengths
+        - **inputs** (batch, time, dim): Tensor containing input vector
+        - **input_lengths** (batch): list of sequence input lengths
+
+    Returns: outputs, output_lengths
+        - **outputs** (batch, out_channels, time): Tensor produces by conformer encoder.
+        - **output_lengths** (batch): list of sequence output lengths
+    """
+    def __init__(
+            self,
+            input_dim: int = 80,
+            encoder_dim: int = 512,
+            num_layers: int = 17,
+            num_attention_heads: int = 8,
+            feed_forward_expansion_factor: int = 4,
+            conv_expansion_factor: int = 2,
+            input_dropout_p: float = 0.1,
+            feed_forward_dropout_p: float = 0.1,
+            attention_dropout_p: float = 0.1,
+            conv_dropout_p: float = 0.1,
+            conv_kernel_size: int = 31,
+            half_step_residual: bool = True,
+    ):
+        super(ConformerEncoder, self).__init__()
+        self.conv_subsample = Conv2dSubampling(in_channels=1, out_channels=encoder_dim)
+        self.input_projection = nn.Sequential(
+            Linear(encoder_dim * (((input_dim - 1) // 2 - 1) // 2), encoder_dim),
+            nn.Dropout(p=input_dropout_p),
+        )
+        self.layers = nn.ModuleList([ConformerBlock(
+            encoder_dim=encoder_dim,
+            num_attention_heads=num_attention_heads,
+            feed_forward_expansion_factor=feed_forward_expansion_factor,
+            conv_expansion_factor=conv_expansion_factor,
+            feed_forward_dropout_p=feed_forward_dropout_p,
+            attention_dropout_p=attention_dropout_p,
+            conv_dropout_p=conv_dropout_p,
+            conv_kernel_size=conv_kernel_size,
+            half_step_residual=half_step_residual,
+        ) for _ in range(num_layers)])
+
+    def count_parameters(self) -> int:
+        """ Count parameters of encoder """
+        return sum([p.numel() for p in self.parameters()])
+
+    def update_dropout(self, dropout_p: float) -> None:
+        """ Update dropout probability of encoder """
+        for name, child in self.named_children():
+            if isinstance(child, nn.Dropout):
+                child.p = dropout_p
+
+    def forward(self, inputs: Tensor, input_lengths: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        Forward propagate a `inputs` for  encoder training.
+
+        Args:
+            inputs (torch.FloatTensor): A input sequence passed to encoder. Typically for inputs this will be a padded
+                `FloatTensor` of size ``(batch, seq_length, dimension)``.
+            input_lengths (torch.LongTensor): The length of input tensor. ``(batch)``
+
+        Returns:
+            (Tensor, Tensor)
+
+            * outputs (torch.FloatTensor): A output sequence of encoder. `FloatTensor` of size
+                ``(batch, seq_length, dimension)``
+            * output_lengths (torch.LongTensor): The length of output tensor. ``(batch)``
+        """
+        outputs, output_lengths = self.conv_subsample(inputs, input_lengths)
+        outputs = self.input_projection(outputs)
+
+        for layer in self.layers:
+            outputs = layer(outputs)
+
+        return outputs, output_lengths

conformer/conformer/feed_forward.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from .activation import Swish
+from .modules import Linear
+
+
+class FeedForwardModule(nn.Module):
+    """
+    Conformer Feed Forward Module follow pre-norm residual units and apply layer normalization within the residual unit
+    and on the input before the first linear layer. This module also apply Swish activation and dropout, which helps
+    regularizing the network.
+
+    Args:
+        encoder_dim (int): Dimension of conformer encoder
+        expansion_factor (int): Expansion factor of feed forward module.
+        dropout_p (float): Ratio of dropout
+
+    Inputs: inputs
+        - **inputs** (batch, time, dim): Tensor contains input sequences
+
+    Outputs: outputs
+        - **outputs** (batch, time, dim): Tensor produces by feed forward module.
+    """
+    def __init__(
+            self,
+            encoder_dim: int = 512,
+            expansion_factor: int = 4,
+            dropout_p: float = 0.1,
+    ) -> None:
+        super(FeedForwardModule, self).__init__()
+        self.sequential = nn.Sequential(
+            nn.LayerNorm(encoder_dim),
+            Linear(encoder_dim, encoder_dim * expansion_factor, bias=True),
+            Swish(),
+            nn.Dropout(p=dropout_p),
+            Linear(encoder_dim * expansion_factor, encoder_dim, bias=True),
+            nn.Dropout(p=dropout_p),
+        )
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return self.sequential(inputs)

conformer/conformer/model.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Tuple
+
+from .encoder import ConformerEncoder
+from .modules import Linear
+
+
+class Conformer(nn.Module):
+    """
+    Conformer: Convolution-augmented Transformer for Speech Recognition
+    The paper used a one-lstm Transducer decoder, currently still only implemented
+    the conformer encoder shown in the paper.
+
+    Args:
+        num_classes (int): Number of classification classes
+        input_dim (int, optional): Dimension of input vector
+        encoder_dim (int, optional): Dimension of conformer encoder
+        num_encoder_layers (int, optional): Number of conformer blocks
+        num_attention_heads (int, optional): Number of attention heads
+        feed_forward_expansion_factor (int, optional): Expansion factor of feed forward module
+        conv_expansion_factor (int, optional): Expansion factor of conformer convolution module
+        feed_forward_dropout_p (float, optional): Probability of feed forward module dropout
+        attention_dropout_p (float, optional): Probability of attention module dropout
+        conv_dropout_p (float, optional): Probability of conformer convolution module dropout
+        conv_kernel_size (int or tuple, optional): Size of the convolving kernel
+        half_step_residual (bool): Flag indication whether to use half step residual or not
+
+    Inputs: inputs, input_lengths
+        - **inputs** (batch, time, dim): Tensor containing input vector
+        - **input_lengths** (batch): list of sequence input lengths
+
+    Returns: outputs, output_lengths
+        - **outputs** (batch, out_channels, time): Tensor produces by conformer.
+        - **output_lengths** (batch): list of sequence output lengths
+    """
+    def __init__(
+            self,
+            num_classes: int,
+            input_dim: int = 80,
+            encoder_dim: int = 512,
+            num_encoder_layers: int = 17,
+            num_attention_heads: int = 8,
+            feed_forward_expansion_factor: int = 4,
+            conv_expansion_factor: int = 2,
+            input_dropout_p: float = 0.1,
+            feed_forward_dropout_p: float = 0.1,
+            attention_dropout_p: float = 0.1,
+            conv_dropout_p: float = 0.1,
+            conv_kernel_size: int = 31,
+            half_step_residual: bool = True,
+    ) -> None:
+        super(Conformer, self).__init__()
+        self.encoder = ConformerEncoder(
+            input_dim=input_dim,
+            encoder_dim=encoder_dim,
+            num_layers=num_encoder_layers,
+            num_attention_heads=num_attention_heads,
+            feed_forward_expansion_factor=feed_forward_expansion_factor,
+            conv_expansion_factor=conv_expansion_factor,
+            input_dropout_p=input_dropout_p,
+            feed_forward_dropout_p=feed_forward_dropout_p,
+            attention_dropout_p=attention_dropout_p,
+            conv_dropout_p=conv_dropout_p,
+            conv_kernel_size=conv_kernel_size,
+            half_step_residual=half_step_residual,
+        )
+        self.fc = Linear(encoder_dim, num_classes, bias=False)
+
+    def count_parameters(self) -> int:
+        """ Count parameters of encoder """
+        return self.encoder.count_parameters()
+
+    def update_dropout(self, dropout_p) -> None:
+        """ Update dropout probability of model """
+        self.encoder.update_dropout(dropout_p)
+
+    def forward(self, inputs: Tensor, input_lengths: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        Forward propagate a `inputs` and `targets` pair for training.
+
+        Args:
+            inputs (torch.FloatTensor): A input sequence passed to encoder. Typically for inputs this will be a padded
+                `FloatTensor` of size ``(batch, seq_length, dimension)``.
+            input_lengths (torch.LongTensor): The length of input tensor. ``(batch)``
+
+        Returns:
+            * predictions (torch.FloatTensor): Result of model predictions.
+        """
+        encoder_outputs, encoder_output_lengths = self.encoder(inputs, input_lengths)
+        outputs = self.fc(encoder_outputs)
+        outputs = nn.functional.log_softmax(outputs, dim=-1)
+        return outputs, encoder_output_lengths

conformer/conformer/model_def.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Tuple
+
+from .encoder import ConformerEncoder
+from .modules import Linear
+
+
+class Conformer(nn.Module):
+    """
+    Conformer: Convolution-augmented Transformer for Speech Recognition
+    The paper used a one-lstm Transducer decoder, currently still only implemented
+    the conformer encoder shown in the paper.
+
+    Args:
+        num_classes (int): Number of classification classes
+        input_dim (int, optional): Dimension of input vector
+        encoder_dim (int, optional): Dimension of conformer encoder
+        num_encoder_layers (int, optional): Number of conformer blocks
+        num_attention_heads (int, optional): Number of attention heads
+        feed_forward_expansion_factor (int, optional): Expansion factor of feed forward module
+        conv_expansion_factor (int, optional): Expansion factor of conformer convolution module
+        feed_forward_dropout_p (float, optional): Probability of feed forward module dropout
+        attention_dropout_p (float, optional): Probability of attention module dropout
+        conv_dropout_p (float, optional): Probability of conformer convolution module dropout
+        conv_kernel_size (int or tuple, optional): Size of the convolving kernel
+        half_step_residual (bool): Flag indication whether to use half step residual or not
+
+    Inputs: inputs, input_lengths
+        - **inputs** (batch, time, dim): Tensor containing input vector
+        - **input_lengths** (batch): list of sequence input lengths
+
+    Returns: outputs, output_lengths
+        - **outputs** (batch, out_channels, time): Tensor produces by conformer.
+        - **output_lengths** (batch): list of sequence output lengths
+    """
+    def __init__(
+            self,
+            input_dim: int = 80,
+            encoder_dim: int = 512,
+            num_encoder_layers: int = 17,
+            num_attention_heads: int = 8,
+            feed_forward_expansion_factor: int = 4,
+            conv_expansion_factor: int = 2,
+            input_dropout_p: float = 0.1,
+            feed_forward_dropout_p: float = 0.1,
+            attention_dropout_p: float = 0.1,
+            conv_dropout_p: float = 0.1,
+            conv_kernel_size: int = 31,
+            half_step_residual: bool = True,
+    ) -> None:
+        super(Conformer, self).__init__()
+        self.encoder = ConformerEncoder(
+            input_dim=input_dim,
+            encoder_dim=encoder_dim,
+            num_layers=num_encoder_layers,
+            num_attention_heads=num_attention_heads,
+            feed_forward_expansion_factor=feed_forward_expansion_factor,
+            conv_expansion_factor=conv_expansion_factor,
+            input_dropout_p=input_dropout_p,
+            feed_forward_dropout_p=feed_forward_dropout_p,
+            attention_dropout_p=attention_dropout_p,
+            conv_dropout_p=conv_dropout_p,
+            conv_kernel_size=conv_kernel_size,
+            half_step_residual=half_step_residual,
+        )
+
+    def count_parameters(self) -> int:
+        """ Count parameters of encoder """
+        return self.encoder.count_parameters()
+
+    def update_dropout(self, dropout_p) -> None:
+        """ Update dropout probability of model """
+        self.encoder.update_dropout(dropout_p)
+
+    def forward(self, inputs: Tensor, input_lengths: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        Forward propagate a `inputs` and `targets` pair for training.
+
+        Args:
+            inputs (torch.FloatTensor): A input sequence passed to encoder. Typically for inputs this will be a padded
+                `FloatTensor` of size ``(batch, seq_length, dimension)``.
+            input_lengths (torch.LongTensor): The length of input tensor. ``(batch)``
+
+        Returns:
+            * predictions (torch.FloatTensor): Result of model predictions.
+        """
+        encoder_outputs, encoder_output_lengths = self.encoder(inputs, input_lengths)
+        return encoder_outputs, encoder_output_lengths

conformer/conformer/modules.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+from torch import Tensor
+
+
+class ResidualConnectionModule(nn.Module):
+    """
+    Residual Connection Module.
+    outputs = (module(inputs) x module_factor + inputs x input_factor)
+    """
+    def __init__(self, module: nn.Module, module_factor: float = 1.0, input_factor: float = 1.0):
+        super(ResidualConnectionModule, self).__init__()
+        self.module = module
+        self.module_factor = module_factor
+        self.input_factor = input_factor
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return (self.module(inputs) * self.module_factor) + (inputs * self.input_factor)
+
+
+class Linear(nn.Module):
+    """
+    Wrapper class of torch.nn.Linear
+    Weight initialize by xavier initialization and bias initialize to zeros.
+    """
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super(Linear, self).__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        init.xavier_uniform_(self.linear.weight)
+        if bias:
+            init.zeros_(self.linear.bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.linear(x)
+
+
+class View(nn.Module):
+    """ Wrapper class of torch.view() for Sequential module. """
+    def __init__(self, shape: tuple, contiguous: bool = False):
+        super(View, self).__init__()
+        self.shape = shape
+        self.contiguous = contiguous
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.contiguous:
+            x = x.contiguous()
+
+        return x.view(*self.shape)
+
+
+class Transpose(nn.Module):
+    """ Wrapper class of torch.transpose() for Sequential module. """
+    def __init__(self, shape: tuple):
+        super(Transpose, self).__init__()
+        self.shape = shape
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.transpose(*self.shape)

model_pinyin.py

@@ -0,0 +1,299 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from torch.optim import Adam
+from torch.optim.lr_scheduler import CosineAnnealingLR
+import math
+from conformer.conformer.model_def import Conformer
+
+class PositionalEncoding(nn.Module):
+    """位置编码模块"""
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)]
+
+
+class CrossAttention(nn.Module):
+    """交叉注意力模块 - 一层交叉注意力加一层自注意力"""
+    def __init__(self, query_dim, key_dim, heads=4, dropout=0.1):
+        super(CrossAttention, self).__init__()
+        # 交叉注意力层
+        self.cross_attn = nn.MultiheadAttention(query_dim, heads, dropout=dropout, batch_first=True)
+        self.norm1 = nn.LayerNorm(query_dim)
+        self.norm2 = nn.LayerNorm(query_dim)
+        self.ffn1 = nn.Sequential(
+            nn.Linear(query_dim, query_dim * 4),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(query_dim * 4, query_dim),
+            nn.Dropout(dropout)
+        )
+        
+        # 自注意力层
+        self.self_attn = nn.MultiheadAttention(query_dim, heads, dropout=dropout, batch_first=True)
+        self.norm3 = nn.LayerNorm(query_dim)
+        self.norm4 = nn.LayerNorm(query_dim)
+        self.ffn2 = nn.Sequential(
+            nn.Linear(query_dim, query_dim * 4),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(query_dim * 4, query_dim),
+            nn.Dropout(dropout)
+        )
+        
+        # 投影层
+        self.proj_key = nn.Linear(key_dim, query_dim)
+        self.proj_value = nn.Linear(key_dim, query_dim)
+        
+    def forward(self, query, key, value, key_padding_mask=None):
+        # 有音频输入时，先进行交叉注意力
+        # 投影key和value到query的维度
+        key_proj = self.proj_key(key)
+        value_proj = self.proj_value(value)
+        
+        # 交叉注意力
+        query_norm = self.norm1(query)
+        cross_attn_output, _ = self.cross_attn(query_norm, key_proj, value_proj, 
+                                             key_padding_mask=key_padding_mask)
+        query = query + cross_attn_output
+        query = query + self.ffn1(self.norm2(query))
+        
+        # 然后进行自注意力
+        query_norm = self.norm3(query)
+        self_attn_output, _ = self.self_attn(query_norm, query_norm, query_norm)
+        query = query + self_attn_output
+        query = query + self.ffn2(self.norm4(query))
+        
+        return query
+    
+
+class MMKWS2(nn.Module):
+    def __init__(
+        self,
+        # anchor
+        text_dim=64,
+        audio_dim=1024,
+        hidden_dim=128,
+        # compare
+        dim=80,
+        encoder_dim=128,
+        num_encoder_layers=6,
+        num_attention_heads=4,
+        dropout=0.1,
+        num_transformer_layers=2
+    ):
+        super(MMKWS2, self).__init__()
+        # 音频嵌入降维
+        self.audio_proj = nn.Linear(audio_dim, hidden_dim)
+        # 文本嵌入投影
+        self.text_proj = nn.Embedding(num_embeddings=402, embedding_dim=hidden_dim) # 401 + padding -1
+        # 位置编码
+        self.pos_enc = PositionalEncoding(hidden_dim)
+        # 交叉注意力模块
+        self.cross_attn = CrossAttention(hidden_dim, hidden_dim, heads=num_attention_heads, dropout=dropout)
+        
+        # Conformer层
+        self.conformer = Conformer(
+            input_dim=dim,
+            encoder_dim=encoder_dim,
+            num_encoder_layers=num_encoder_layers,
+            num_attention_heads=num_attention_heads,
+        )
+        
+        # 特征映射层（将conformer输出维度映射到hidden_dim）
+        self.feat_proj = nn.Linear(encoder_dim, hidden_dim)
+        
+        # Transformer层
+        self.transformer_encoder = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=hidden_dim,
+                nhead=num_attention_heads,
+                dim_feedforward=hidden_dim*4,
+                dropout=dropout,
+                batch_first=True
+            ),
+            num_layers=num_transformer_layers
+        )
+        
+        # GRU分类器
+        self.gru = nn.GRU(hidden_dim, hidden_dim, batch_first=True, bidirectional=True)
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_dim*2, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, 1)
+        )
+        
+        # 序列标签预测
+        self.seq_classifier = nn.Linear(hidden_dim, 1)
+        
+    def forward(self, anchor_wave_embedding, anchor_text_embedding, compare_wave, compare_lengths):
+        batch_size = anchor_wave_embedding.size(0)
+        
+        # 1. 处理anchor_text嵌入
+        text_feat = self.text_proj(anchor_text_embedding)  # [B, S, hidden_dim]
+        text_feat = self.pos_enc(text_feat)
+        
+        # 2. 处理anchor_wave音频嵌入
+        audio_feat = self.audio_proj(anchor_wave_embedding)  # [B, S, hidden_dim]
+        audio_feat = self.pos_enc(audio_feat)
+        
+        # 3. 交叉注意力：文本和音频特征融合
+        fused_feat = self.cross_attn(text_feat, audio_feat, audio_feat, key_padding_mask=None)  # [B, S, hidden_dim]
+        
+        # 4. 处理compare_wave的fbank特征
+        compare_feat = self.conformer(compare_wave, compare_lengths)[0]  # [B, T, encoder_dim]        
+        compare_feat = self.feat_proj(compare_feat)  # [B, T, hidden_dim]
+        compare_feat = self.pos_enc(compare_feat)
+
+        # 5. 合并特征
+        text_len = fused_feat.size(1)
+        combined_feat = torch.cat([fused_feat, compare_feat], dim=1)  # [B, S+T, hidden_dim]
+        combined_feat = self.transformer_encoder(combined_feat) # [B, S+T, hidden_dim]
+            
+        # 7. GRU分类
+        gru_out, _ = self.gru(combined_feat)  # [B, S+T, hidden_dim*2]
+        
+        # 全局分类
+        global_feat = gru_out[:, -1, :]  # 取最后一个时间步
+        logits = self.classifier(global_feat).squeeze(-1)  # [B]
+        
+        # 序列标签预测
+        seq_logits = self.seq_classifier(combined_feat[:, :text_len, :]).squeeze(-1)  # [B, S]
+        return logits, seq_logits
+    
+    
+    def enrollment(self, anchor_wave_embedding, anchor_text_embedding):
+        batch_size = anchor_wave_embedding.size(0)
+        
+        # 1. 处理anchor_text嵌入
+        text_feat = self.text_proj(anchor_text_embedding)  # [B, S, hidden_dim]
+        text_feat = self.pos_enc(text_feat)
+        
+        # 2. 处理anchor_wave音频嵌入
+        audio_feat = self.audio_proj(anchor_wave_embedding)  # [B, S, hidden_dim]
+        audio_feat = self.pos_enc(audio_feat)
+        
+        # 3. 交叉注意力：文本和音频特征融合
+        fused_feat = self.cross_attn(text_feat, audio_feat, audio_feat, key_padding_mask=None)  # [B, S, hidden_dim]
+
+        return fused_feat
+    
+    def verification(self, fused_feat, compare_wave, compare_lengths):
+        batch_size = fused_feat.size(0)
+
+        # 4. 处理compare_wave的fbank特征
+        compare_feat = self.conformer(compare_wave, compare_lengths)[0]  # [B, T, encoder_dim]        
+        compare_feat = self.feat_proj(compare_feat)  # [B, T, hidden_dim]
+        compare_feat = self.pos_enc(compare_feat)
+
+        # 5. 合并特征
+        text_len = fused_feat.size(1)
+        combined_feat = torch.cat([fused_feat, compare_feat], dim=1)  # [B, S+T, hidden_dim]
+        combined_feat = self.transformer_encoder(combined_feat) # [B, S+T, hidden_dim]
+            
+        # 7. GRU分类
+        gru_out, _ = self.gru(combined_feat)  # [B, S+T, hidden_dim*2]
+        
+        # 全局分类
+        global_feat = gru_out[:, -1, :]  # 取最后一个时间步
+        logits = self.classifier(global_feat).squeeze(-1)  # [B]
+        
+        return logits
+
+
+def count_verification_params(model):
+    modules = [
+        model.conformer,
+        model.feat_proj,
+        model.transformer_encoder,
+        model.gru,
+        model.classifier
+    ]
+    total = 0
+    for m in modules:
+        total += sum(p.numel() for p in m.parameters())
+    return total
+
+model = MMKWS2(
+    text_dim=64,
+    audio_dim=1024,
+    hidden_dim=128,
+    dim=80,
+    encoder_dim=128,
+    num_encoder_layers=6,
+    num_attention_heads=4,
+    dropout=0.1,
+    num_transformer_layers=2
+)
+print(f"verification相关参数量: {count_verification_params(model):,}") # 3.5M模型参数量
+
+# if __name__ == "__main__":
+#     # 创建一个示例batch
+#     batch_size = 2
+    
+#     # 创建模拟数据
+#     anchor_embedding = torch.randn(batch_size, 8, 64)  # 文本嵌入
+#     anchor_wave = torch.randn(batch_size, 256, 1024)    # 音频嵌入
+#     compare_wave = torch.randn(batch_size, 45, 80)   # Fbank特征
+    
+#     # 创建长度信息
+#     anchor_lengths = torch.LongTensor([8, 6])  # 两个样本的实际长度
+#     compare_lengths = torch.LongTensor([45, 40])
+    
+#     # 创建模型
+#     model = MMKWS2(
+#         text_dim=64,
+#         audio_dim=1024,
+#         hidden_dim=128,
+#         dim=80,
+#         encoder_dim=128,
+#         num_encoder_layers=6,
+#         num_attention_heads=4,
+#         dropout=0.1,
+#         num_transformer_layers=2
+#     )
+    
+#     # 计算模型参数量
+#     total_params = sum(p.numel() for p in model.parameters())
+#     trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    
+#     print(f"模型总参数量: {total_params:,}")
+#     print(f"可训练参数量: {trainable_params:,}")
+    
+#     # 打印模型结构
+#     print("\n模型结构:")
+#     print(model)
+    
+#     # 模型推理
+#     print("\n开始推理...")
+#     model.eval()
+#     with torch.no_grad():
+        
+#         print(anchor_embedding.shape)
+#         print(anchor_wave.shape)
+#         print(compare_wave.shape)
+#         print(anchor_lengths.shape)
+#         print(compare_lengths.shape)
+#         # 完整输入推理
+#         logits, seq_logits, text_len = model(
+#             anchor_embedding=anchor_embedding,
+#             anchor_wave=anchor_wave,
+#             compare_wave=compare_wave,
+#             anchor_lengths=anchor_lengths,
+#             compare_lengths=compare_lengths
+#         )
+        
+#         print("\n推理结果:")
+#         print(f"分类logits形状: {logits.shape}")
+#         print(f"序列logits形状: {seq_logits.shape}")

pinyin_index.json

@@ -0,0 +1,808 @@
+{
+    "pinyin_to_index": {
+      "a": 0,
+      "ai": 1,
+      "an": 2,
+      "ang": 3,
+      "ao": 4,
+      "ba": 5,
+      "bai": 6,
+      "ban": 7,
+      "bang": 8,
+      "bao": 9,
+      "bei": 10,
+      "ben": 11,
+      "beng": 12,
+      "bi": 13,
+      "bian": 14,
+      "biao": 15,
+      "bie": 16,
+      "bin": 17,
+      "bing": 18,
+      "bo": 19,
+      "bu": 20,
+      "ca": 21,
+      "cai": 22,
+      "can": 23,
+      "cang": 24,
+      "cao": 25,
+      "ce": 26,
+      "cen": 27,
+      "ceng": 28,
+      "cha": 29,
+      "chai": 30,
+      "chan": 31,
+      "chang": 32,
+      "chao": 33,
+      "che": 34,
+      "chen": 35,
+      "cheng": 36,
+      "chi": 37,
+      "chong": 38,
+      "chou": 39,
+      "chu": 40,
+      "chuai": 41,
+      "chuan": 42,
+      "chuang": 43,
+      "chui": 44,
+      "chun": 45,
+      "chuo": 46,
+      "ci": 47,
+      "cong": 48,
+      "cou": 49,
+      "cu": 50,
+      "cuan": 51,
+      "cui": 52,
+      "cun": 53,
+      "cuo": 54,
+      "da": 55,
+      "dai": 56,
+      "dan": 57,
+      "dang": 58,
+      "dao": 59,
+      "de": 60,
+      "dei": 61,
+      "deng": 62,
+      "di": 63,
+      "dia": 64,
+      "dian": 65,
+      "diao": 66,
+      "die": 67,
+      "ding": 68,
+      "diu": 69,
+      "dong": 70,
+      "dou": 71,
+      "du": 72,
+      "duan": 73,
+      "dui": 74,
+      "dun": 75,
+      "duo": 76,
+      "e": 77,
+      "en": 78,
+      "er": 79,
+      "fa": 80,
+      "fan": 81,
+      "fang": 82,
+      "fei": 83,
+      "fen": 84,
+      "feng": 85,
+      "fo": 86,
+      "fou": 87,
+      "fu": 88,
+      "ga": 89,
+      "gai": 90,
+      "gan": 91,
+      "gang": 92,
+      "gao": 93,
+      "ge": 94,
+      "gei": 95,
+      "gen": 96,
+      "geng": 97,
+      "gong": 98,
+      "gou": 99,
+      "gu": 100,
+      "gua": 101,
+      "guai": 102,
+      "guan": 103,
+      "guang": 104,
+      "gui": 105,
+      "gun": 106,
+      "guo": 107,
+      "ha": 108,
+      "hai": 109,
+      "han": 110,
+      "hang": 111,
+      "hao": 112,
+      "he": 113,
+      "hei": 114,
+      "hen": 115,
+      "heng": 116,
+      "hong": 117,
+      "hou": 118,
+      "hu": 119,
+      "hua": 120,
+      "huai": 121,
+      "huan": 122,
+      "huang": 123,
+      "hui": 124,
+      "hun": 125,
+      "huo": 126,
+      "ji": 127,
+      "jia": 128,
+      "jian": 129,
+      "jiang": 130,
+      "jiao": 131,
+      "jie": 132,
+      "jin": 133,
+      "jing": 134,
+      "jiong": 135,
+      "jiu": 136,
+      "ju": 137,
+      "juan": 138,
+      "jue": 139,
+      "jun": 140,
+      "ka": 141,
+      "kai": 142,
+      "kan": 143,
+      "kang": 144,
+      "kao": 145,
+      "ke": 146,
+      "ken": 147,
+      "keng": 148,
+      "kong": 149,
+      "kou": 150,
+      "ku": 151,
+      "kua": 152,
+      "kuai": 153,
+      "kuan": 154,
+      "kuang": 155,
+      "kui": 156,
+      "kun": 157,
+      "kuo": 158,
+      "la": 159,
+      "lai": 160,
+      "lan": 161,
+      "lang": 162,
+      "lao": 163,
+      "le": 164,
+      "lei": 165,
+      "leng": 166,
+      "li": 167,
+      "lia": 168,
+      "lian": 169,
+      "liang": 170,
+      "liao": 171,
+      "lie": 172,
+      "lin": 173,
+      "ling": 174,
+      "liu": 175,
+      "long": 176,
+      "lou": 177,
+      "lu": 178,
+      "luan": 179,
+      "lun": 180,
+      "luo": 181,
+      "lv": 182,
+      "lve": 183,
+      "ma": 184,
+      "mai": 185,
+      "man": 186,
+      "mang": 187,
+      "mao": 188,
+      "me": 189,
+      "mei": 190,
+      "men": 191,
+      "meng": 192,
+      "mi": 193,
+      "mian": 194,
+      "miao": 195,
+      "mie": 196,
+      "min": 197,
+      "ming": 198,
+      "miu": 199,
+      "mo": 200,
+      "mou": 201,
+      "mu": 202,
+      "na": 203,
+      "nai": 204,
+      "nan": 205,
+      "nang": 206,
+      "nao": 207,
+      "ne": 208,
+      "nei": 209,
+      "nen": 210,
+      "neng": 211,
+      "ni": 212,
+      "nian": 213,
+      "niang": 214,
+      "niao": 215,
+      "nie": 216,
+      "nin": 217,
+      "ning": 218,
+      "niu": 219,
+      "nong": 220,
+      "nu": 221,
+      "nuan": 222,
+      "nuo": 223,
+      "nv": 224,
+      "nve": 225,
+      "o": 226,
+      "ou": 227,
+      "pa": 228,
+      "pai": 229,
+      "pan": 230,
+      "pang": 231,
+      "pao": 232,
+      "pei": 233,
+      "pen": 234,
+      "peng": 235,
+      "pi": 236,
+      "pian": 237,
+      "piao": 238,
+      "pie": 239,
+      "pin": 240,
+      "ping": 241,
+      "po": 242,
+      "pou": 243,
+      "pu": 244,
+      "qi": 245,
+      "qia": 246,
+      "qian": 247,
+      "qiang": 248,
+      "qiao": 249,
+      "qie": 250,
+      "qin": 251,
+      "qing": 252,
+      "qiong": 253,
+      "qiu": 254,
+      "qu": 255,
+      "quan": 256,
+      "que": 257,
+      "qun": 258,
+      "ran": 259,
+      "rang": 260,
+      "rao": 261,
+      "re": 262,
+      "ren": 263,
+      "reng": 264,
+      "ri": 265,
+      "rong": 266,
+      "rou": 267,
+      "ru": 268,
+      "ruan": 269,
+      "rui": 270,
+      "run": 271,
+      "ruo": 272,
+      "sa": 273,
+      "sai": 274,
+      "san": 275,
+      "sang": 276,
+      "sao": 277,
+      "se": 278,
+      "sen": 279,
+      "seng": 280,
+      "sha": 281,
+      "shai": 282,
+      "shan": 283,
+      "shang": 284,
+      "shao": 285,
+      "she": 286,
+      "shei": 287,
+      "shen": 288,
+      "sheng": 289,
+      "shi": 290,
+      "shou": 291,
+      "shu": 292,
+      "shua": 293,
+      "shuai": 294,
+      "shuan": 295,
+      "shuang": 296,
+      "shui": 297,
+      "shun": 298,
+      "shuo": 299,
+      "si": 300,
+      "song": 301,
+      "sou": 302,
+      "su": 303,
+      "suan": 304,
+      "sui": 305,
+      "sun": 306,
+      "suo": 307,
+      "ta": 308,
+      "tai": 309,
+      "tan": 310,
+      "tang": 311,
+      "tao": 312,
+      "te": 313,
+      "teng": 314,
+      "ti": 315,
+      "tian": 316,
+      "tiao": 317,
+      "tie": 318,
+      "ting": 319,
+      "tong": 320,
+      "tou": 321,
+      "tu": 322,
+      "tuan": 323,
+      "tui": 324,
+      "tun": 325,
+      "tuo": 326,
+      "wa": 327,
+      "wai": 328,
+      "wan": 329,
+      "wang": 330,
+      "wei": 331,
+      "wen": 332,
+      "weng": 333,
+      "wo": 334,
+      "wu": 335,
+      "xi": 336,
+      "xia": 337,
+      "xian": 338,
+      "xiang": 339,
+      "xiao": 340,
+      "xie": 341,
+      "xin": 342,
+      "xing": 343,
+      "xiong": 344,
+      "xiu": 345,
+      "xu": 346,
+      "xuan": 347,
+      "xue": 348,
+      "xun": 349,
+      "ya": 350,
+      "yan": 351,
+      "yang": 352,
+      "yao": 353,
+      "ye": 354,
+      "yi": 355,
+      "yin": 356,
+      "ying": 357,
+      "yo": 358,
+      "yong": 359,
+      "you": 360,
+      "yu": 361,
+      "yuan": 362,
+      "yue": 363,
+      "yun": 364,
+      "za": 365,
+      "zai": 366,
+      "zan": 367,
+      "zang": 368,
+      "zao": 369,
+      "ze": 370,
+      "zei": 371,
+      "zen": 372,
+      "zeng": 373,
+      "zha": 374,
+      "zhai": 375,
+      "zhan": 376,
+      "zhang": 377,
+      "zhao": 378,
+      "zhe": 379,
+      "zhen": 380,
+      "zheng": 381,
+      "zhi": 382,
+      "zhong": 383,
+      "zhou": 384,
+      "zhu": 385,
+      "zhua": 386,
+      "zhuai": 387,
+      "zhuan": 388,
+      "zhuang": 389,
+      "zhui": 390,
+      "zhun": 391,
+      "zhuo": 392,
+      "zi": 393,
+      "zong": 394,
+      "zou": 395,
+      "zu": 396,
+      "zuan": 397,
+      "zui": 398,
+      "zun": 399,
+      "zuo": 400
+    },
+    "index_to_pinyin": {
+      "0": "a",
+      "1": "ai",
+      "2": "an",
+      "3": "ang",
+      "4": "ao",
+      "5": "ba",
+      "6": "bai",
+      "7": "ban",
+      "8": "bang",
+      "9": "bao",
+      "10": "bei",
+      "11": "ben",
+      "12": "beng",
+      "13": "bi",
+      "14": "bian",
+      "15": "biao",
+      "16": "bie",
+      "17": "bin",
+      "18": "bing",
+      "19": "bo",
+      "20": "bu",
+      "21": "ca",
+      "22": "cai",
+      "23": "can",
+      "24": "cang",
+      "25": "cao",
+      "26": "ce",
+      "27": "cen",
+      "28": "ceng",
+      "29": "cha",
+      "30": "chai",
+      "31": "chan",
+      "32": "chang",
+      "33": "chao",
+      "34": "che",
+      "35": "chen",
+      "36": "cheng",
+      "37": "chi",
+      "38": "chong",
+      "39": "chou",
+      "40": "chu",
+      "41": "chuai",
+      "42": "chuan",
+      "43": "chuang",
+      "44": "chui",
+      "45": "chun",
+      "46": "chuo",
+      "47": "ci",
+      "48": "cong",
+      "49": "cou",
+      "50": "cu",
+      "51": "cuan",
+      "52": "cui",
+      "53": "cun",
+      "54": "cuo",
+      "55": "da",
+      "56": "dai",
+      "57": "dan",
+      "58": "dang",
+      "59": "dao",
+      "60": "de",
+      "61": "dei",
+      "62": "deng",
+      "63": "di",
+      "64": "dia",
+      "65": "dian",
+      "66": "diao",
+      "67": "die",
+      "68": "ding",
+      "69": "diu",
+      "70": "dong",
+      "71": "dou",
+      "72": "du",
+      "73": "duan",
+      "74": "dui",
+      "75": "dun",
+      "76": "duo",
+      "77": "e",
+      "78": "en",
+      "79": "er",
+      "80": "fa",
+      "81": "fan",
+      "82": "fang",
+      "83": "fei",
+      "84": "fen",
+      "85": "feng",
+      "86": "fo",
+      "87": "fou",
+      "88": "fu",
+      "89": "ga",
+      "90": "gai",
+      "91": "gan",
+      "92": "gang",
+      "93": "gao",
+      "94": "ge",
+      "95": "gei",
+      "96": "gen",
+      "97": "geng",
+      "98": "gong",
+      "99": "gou",
+      "100": "gu",
+      "101": "gua",
+      "102": "guai",
+      "103": "guan",
+      "104": "guang",
+      "105": "gui",
+      "106": "gun",
+      "107": "guo",
+      "108": "ha",
+      "109": "hai",
+      "110": "han",
+      "111": "hang",
+      "112": "hao",
+      "113": "he",
+      "114": "hei",
+      "115": "hen",
+      "116": "heng",
+      "117": "hong",
+      "118": "hou",
+      "119": "hu",
+      "120": "hua",
+      "121": "huai",
+      "122": "huan",
+      "123": "huang",
+      "124": "hui",
+      "125": "hun",
+      "126": "huo",
+      "127": "ji",
+      "128": "jia",
+      "129": "jian",
+      "130": "jiang",
+      "131": "jiao",
+      "132": "jie",
+      "133": "jin",
+      "134": "jing",
+      "135": "jiong",
+      "136": "jiu",
+      "137": "ju",
+      "138": "juan",
+      "139": "jue",
+      "140": "jun",
+      "141": "ka",
+      "142": "kai",
+      "143": "kan",
+      "144": "kang",
+      "145": "kao",
+      "146": "ke",
+      "147": "ken",
+      "148": "keng",
+      "149": "kong",
+      "150": "kou",
+      "151": "ku",
+      "152": "kua",
+      "153": "kuai",
+      "154": "kuan",
+      "155": "kuang",
+      "156": "kui",
+      "157": "kun",
+      "158": "kuo",
+      "159": "la",
+      "160": "lai",
+      "161": "lan",
+      "162": "lang",
+      "163": "lao",
+      "164": "le",
+      "165": "lei",
+      "166": "leng",
+      "167": "li",
+      "168": "lia",
+      "169": "lian",
+      "170": "liang",
+      "171": "liao",
+      "172": "lie",
+      "173": "lin",
+      "174": "ling",
+      "175": "liu",
+      "176": "long",
+      "177": "lou",
+      "178": "lu",
+      "179": "luan",
+      "180": "lun",
+      "181": "luo",
+      "182": "lv",
+      "183": "lve",
+      "184": "ma",
+      "185": "mai",
+      "186": "man",
+      "187": "mang",
+      "188": "mao",
+      "189": "me",
+      "190": "mei",
+      "191": "men",
+      "192": "meng",
+      "193": "mi",
+      "194": "mian",
+      "195": "miao",
+      "196": "mie",
+      "197": "min",
+      "198": "ming",
+      "199": "miu",
+      "200": "mo",
+      "201": "mou",
+      "202": "mu",
+      "203": "na",
+      "204": "nai",
+      "205": "nan",
+      "206": "nang",
+      "207": "nao",
+      "208": "ne",
+      "209": "nei",
+      "210": "nen",
+      "211": "neng",
+      "212": "ni",
+      "213": "nian",
+      "214": "niang",
+      "215": "niao",
+      "216": "nie",
+      "217": "nin",
+      "218": "ning",
+      "219": "niu",
+      "220": "nong",
+      "221": "nu",
+      "222": "nuan",
+      "223": "nuo",
+      "224": "nv",
+      "225": "nve",
+      "226": "o",
+      "227": "ou",
+      "228": "pa",
+      "229": "pai",
+      "230": "pan",
+      "231": "pang",
+      "232": "pao",
+      "233": "pei",
+      "234": "pen",
+      "235": "peng",
+      "236": "pi",
+      "237": "pian",
+      "238": "piao",
+      "239": "pie",
+      "240": "pin",
+      "241": "ping",
+      "242": "po",
+      "243": "pou",
+      "244": "pu",
+      "245": "qi",
+      "246": "qia",
+      "247": "qian",
+      "248": "qiang",
+      "249": "qiao",
+      "250": "qie",
+      "251": "qin",
+      "252": "qing",
+      "253": "qiong",
+      "254": "qiu",
+      "255": "qu",
+      "256": "quan",
+      "257": "que",
+      "258": "qun",
+      "259": "ran",
+      "260": "rang",
+      "261": "rao",
+      "262": "re",
+      "263": "ren",
+      "264": "reng",
+      "265": "ri",
+      "266": "rong",
+      "267": "rou",
+      "268": "ru",
+      "269": "ruan",
+      "270": "rui",
+      "271": "run",
+      "272": "ruo",
+      "273": "sa",
+      "274": "sai",
+      "275": "san",
+      "276": "sang",
+      "277": "sao",
+      "278": "se",
+      "279": "sen",
+      "280": "seng",
+      "281": "sha",
+      "282": "shai",
+      "283": "shan",
+      "284": "shang",
+      "285": "shao",
+      "286": "she",
+      "287": "shei",
+      "288": "shen",
+      "289": "sheng",
+      "290": "shi",
+      "291": "shou",
+      "292": "shu",
+      "293": "shua",
+      "294": "shuai",
+      "295": "shuan",
+      "296": "shuang",
+      "297": "shui",
+      "298": "shun",
+      "299": "shuo",
+      "300": "si",
+      "301": "song",
+      "302": "sou",
+      "303": "su",
+      "304": "suan",
+      "305": "sui",
+      "306": "sun",
+      "307": "suo",
+      "308": "ta",
+      "309": "tai",
+      "310": "tan",
+      "311": "tang",
+      "312": "tao",
+      "313": "te",
+      "314": "teng",
+      "315": "ti",
+      "316": "tian",
+      "317": "tiao",
+      "318": "tie",
+      "319": "ting",
+      "320": "tong",
+      "321": "tou",
+      "322": "tu",
+      "323": "tuan",
+      "324": "tui",
+      "325": "tun",
+      "326": "tuo",
+      "327": "wa",
+      "328": "wai",
+      "329": "wan",
+      "330": "wang",
+      "331": "wei",
+      "332": "wen",
+      "333": "weng",
+      "334": "wo",
+      "335": "wu",
+      "336": "xi",
+      "337": "xia",
+      "338": "xian",
+      "339": "xiang",
+      "340": "xiao",
+      "341": "xie",
+      "342": "xin",
+      "343": "xing",
+      "344": "xiong",
+      "345": "xiu",
+      "346": "xu",
+      "347": "xuan",
+      "348": "xue",
+      "349": "xun",
+      "350": "ya",
+      "351": "yan",
+      "352": "yang",
+      "353": "yao",
+      "354": "ye",
+      "355": "yi",
+      "356": "yin",
+      "357": "ying",
+      "358": "yo",
+      "359": "yong",
+      "360": "you",
+      "361": "yu",
+      "362": "yuan",
+      "363": "yue",
+      "364": "yun",
+      "365": "za",
+      "366": "zai",
+      "367": "zan",
+      "368": "zang",
+      "369": "zao",
+      "370": "ze",
+      "371": "zei",
+      "372": "zen",
+      "373": "zeng",
+      "374": "zha",
+      "375": "zhai",
+      "376": "zhan",
+      "377": "zhang",
+      "378": "zhao",
+      "379": "zhe",
+      "380": "zhen",
+      "381": "zheng",
+      "382": "zhi",
+      "383": "zhong",
+      "384": "zhou",
+      "385": "zhu",
+      "386": "zhua",
+      "387": "zhuai",
+      "388": "zhuan",
+      "389": "zhuang",
+      "390": "zhui",
+      "391": "zhun",
+      "392": "zhuo",
+      "393": "zi",
+      "394": "zong",
+      "395": "zou",
+      "396": "zu",
+      "397": "zuan",
+      "398": "zui",
+      "399": "zun",
+      "400": "zuo"
+    }
+  }

stepstep=024500.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6bd2c480df21e2cd9ee5eadef5d67e69b059e00c30159889502ac1848132950
+size 714212484

train_pinyin.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torch.optim.lr_scheduler import LambdaLR
+
+import numpy as np
+from sklearn.metrics import roc_auc_score
+from sklearn.metrics import roc_curve
+
+import pytorch_lightning as pl
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from transformers import Wav2Vec2FeatureExtractor, HubertModel
+from model_pinyin import MMKWS2
+from torch.utils.data import Dataset, ConcatDataset
+
+class MMKWS2_Wrapper(LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = MMKWS2()
+        self.criterion = nn.BCEWithLogitsLoss()
+        # 将hubert_model设为临时变量而非类属性
+        hubert_model = HubertModel.from_pretrained("TencentGameMate/chinese-hubert-large").half().eval()
+        self._hubert_model = hubert_model  # 使用下划线前缀表示内部使用
+        
+    def training_step(self, batch, batch_idx):
+        anchor_wave, anchor_text_embedding, compare_wave, compare_lengths, label, seq_label = \
+            batch['anchor_wave'], batch['anchor_embedding'], batch['compare_wave'], batch['compare_lengths'], batch['label'], batch['seq_label']
+            
+        with torch.no_grad():
+            outputs = self._hubert_model(anchor_wave.half())
+            anchor_wave_embedding = outputs.last_hidden_state
+        
+        anchor_wave_embedding = anchor_wave_embedding.to(anchor_wave.dtype)
+        
+        logits, seq_logits = self.model(
+            anchor_wave_embedding,
+            anchor_text_embedding,
+            compare_wave,
+            compare_lengths
+        )
+        
+        # 句级二分类loss
+        utt_loss = self.criterion(logits, label.float())
+        
+        # 序列loss（mask掉seq_label为-1的部分）
+        mask = (seq_label != -1).float()
+        seq_label_valid = seq_label.clone()
+        seq_label_valid[seq_label == -1] = 0  # 避免-1影响loss
+        
+        seq_loss = F.binary_cross_entropy_with_logits(
+            seq_logits, seq_label_valid.float(), weight=mask, reduction='sum'
+        ) / (mask.sum() + 1e-6)
+        
+        loss = utt_loss + seq_loss
+        
+        # 每500步记录日志
+        self.log('train/utt_loss', utt_loss, on_step=True, on_epoch=False, prog_bar=True)
+        self.log('train/seq_loss', seq_loss, on_step=True, on_epoch=False, prog_bar=True)
+        self.log('train/loss', loss, on_step=True, on_epoch=False, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        
+        def lr_lambda(step):
+            return 0.95 ** (step // 1000)
+            
+        scheduler = torch.optim.lr_scheduler.LambdaLR(
+            optimizer, 
+            lr_lambda=lr_lambda
+        )
+        
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",  # 按步更新
+                "frequency": 1
+            },
+        }
+
+
+# 3. 设置 Trainer 和训练
+if __name__ == "__main__":
+    pl.seed_everything(2024)
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = '7'
+    from dataset_pinyin import PairDataset
+    from dataloader_pinyin import get_dataloader
+    
+    # 创建数据集
+    dataset1 = PairDataset(
+        '/nvme01/aizq/kws-agent/data/anchor_pairs.parquet', 
+        '/nvme01/aizq/kws-agent/data/WenetPhrase_base/M_S',
+        augment=True
+    )
+    dataset2 = PairDataset(
+        '/nvme01/aizq/kws-agent/data/synthetic_pairs.parquet', 
+        '/nvme01/aizq/kws-agent/data/WenetPhrase_base/M_S',
+        augment=True
+    )
+    dataset = ConcatDataset([dataset1, dataset2])
+    # 创建dataloader
+    dataloader = get_dataloader(dataset, batch_size=1024)
+
+    model = MMKWS2_Wrapper()
+    model_checkpoint = ModelCheckpoint(
+        dirpath="/nvme01/aizq/kws-agent/ckpts",
+        filename="step{step:06d}",
+        save_top_k=-1,
+        save_on_train_epoch_end=False,  # 按训练步数保存
+        every_n_train_steps=500       # 每10k步保存一次
+    )
+    logger = pl.loggers.TensorBoardLogger('/nvme01/aizq/kws-agent/logs/', name='MMKWS+')
+    trainer = Trainer(
+        devices=1, 
+        accelerator='gpu', 
+        logger=logger,
+        max_epochs=4,  # 训练1轮
+        callbacks=[model_checkpoint], 
+        accumulate_grad_batches=2, # 2048 batchsize
+    )
+    trainer.fit(model, train_dataloaders=dataloader)

tts-2025-04-27@197cd135f2a2451b9cab9cf2add1c1ab.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3947a005d2c85f427cead98d3aeb65200bb85a956383aad43467b1b6a7f0c5f1
+size 279404

ui.py

@@ -0,0 +1,167 @@
+import sys
+import os
+import json
+import time
+import numpy as np
+import torch
+import torchaudio
+import gradio as gr
+from torchaudio.compliance.kaldi import fbank
+from pypinyin import lazy_pinyin
+from train_pinyin import MMKWS2_Wrapper
+
+# 设备与模型加载
+# device = torch.device("cuda:4")
+device = torch.device("cpu")
+wrapper = MMKWS2_Wrapper.load_from_checkpoint(
+    "stepstep=024500.ckpt",
+    map_location=device,
+)
+wrapper.eval()
+
+# 注册信息
+registered = {"text": "", "audios": []}
+enroll = None
+enroll_text = None
+last_wake_time = 0
+
+def load_pinyin_index(save_path):
+    """加载拼音索引映射"""
+    with open(save_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    return data["pinyin_to_index"], data["index_to_pinyin"]
+
+pinyin_to_index, index_to_pinyin = load_pinyin_index("pinyin_index.json")
+
+def add_audio(text, audio, audio_list):
+    if not text:
+        return audio_list, "请先输入唤醒词文本"
+    if audio is None or audio[1] is None or len(audio[1]) == 0:
+        return audio_list, "请上传或录制音频"
+    audio_list = audio_list or []
+    if len(audio_list) >= 5:
+        return audio_list, "最多支持5条音频"
+    audio_list.append(audio)
+    return audio_list, f"已录入 {len(audio_list)} 条音频"
+
+def register_keyword(text, audio_list):
+    if not text:
+        return gr.update(value="请先输入唤醒词文本")
+    if not audio_list or len(audio_list) == 0:
+        return gr.update(value="请至少上传或录制一条音频")
+    registered["text"] = text
+    registered["audios"] = audio_list
+    global enroll_text
+    enroll_text = text
+    fused_feats = []
+    for audio in audio_list:
+        anchor_wave, _ = torchaudio.load(audio)
+        anchor_text_embedding = torch.tensor([pinyin_to_index[p] + 1 for p in lazy_pinyin(text)])
+        anchor_wave = anchor_wave.to(device)
+        anchor_text_embedding = anchor_text_embedding.to(device).unsqueeze(0)
+        with torch.no_grad():
+            outputs = wrapper._hubert_model(anchor_wave.half())
+            anchor_wave_embedding = outputs.last_hidden_state
+        anchor_wave_embedding = anchor_wave_embedding.to(anchor_wave.dtype)
+        fused_feat = wrapper.model.enrollment(
+            anchor_wave_embedding,
+            anchor_text_embedding
+        )
+        fused_feats.append(fused_feat)
+    fused_feats = torch.cat(fused_feats, dim=0)
+    fused_feats, _ = fused_feats.max(dim=0)
+    fused_feats = fused_feats.unsqueeze(0)
+    global enroll
+    enroll = fused_feats
+    return gr.update(value=f"注册完成，唤醒词：{text}，音频数：{len(audio_list)}")
+
+def update_gallery(audio_list):
+    if audio_list and len(audio_list) > 0:
+        return gr.update(visible=True, value=audio_list[-1])
+    else:
+        return gr.update(visible=False, value=None)
+
+def streaming_detect_handler(current_audio, state, audio_player):
+    global last_wake_time, enroll_text, enroll
+    if current_audio is None or current_audio[1] is None or len(current_audio[1]) == 0:
+        return state, gr.update()
+    if enroll_text is None:
+        return state, gr.update()
+    pad = len(enroll_text) * 5
+    state = (state or []) + [current_audio]
+    state = state[-pad:]
+    if len(state) < pad:
+        return state, gr.update()
+    sr = state[0][0]
+    audio_list = [x[1] for x in state]
+    audio_concat = np.concatenate(audio_list, axis=0)
+    audio_concat = audio_concat.astype(np.float32) / 32768.0
+    audio_concat = torch.from_numpy(audio_concat).unsqueeze(0)
+    audio_concat = torchaudio.functional.resample(audio_concat, sr, 16000)
+    audio_concat = audio_concat / torch.max(torch.abs(audio_concat))
+    compare_wave = fbank(audio_concat, num_mel_bins=80)
+    compare_wave = compare_wave.to(device).unsqueeze(0)
+    compare_lengths = torch.tensor([compare_wave.size(1)], device=compare_wave.device)
+    if enroll is None:
+        return state, None
+    current_time = time.time()
+    if current_time - last_wake_time <= 2:
+        return state, gr.update()
+    with torch.no_grad():
+        preds = wrapper.model.verification(
+            enroll,
+            compare_wave,
+            compare_lengths
+        )
+    preds = torch.sigmoid(preds).item()
+    if preds >= 0.85:
+        print(f"Wake up! {preds}")
+        last_wake_time = current_time
+        audio_path = "tts-2025-04-27@197cd135f2a2451b9cab9cf2add1c1ab.wav"
+        state = []
+        return state, gr.Audio(value=audio_path, visible=True, autoplay=True)
+    if preds >= 0.6:
+        print(f"Preds! {preds}")
+    return state, None
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 自定义关键词检测 Demo")
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## 注册唤醒词")
+            text_input = gr.Textbox(label="唤醒词文本", placeholder="请输入唤醒词")
+            audio_list_state = gr.State([])
+            audio_input = gr.Audio(label="上传或录制音频", type="filepath")
+            add_btn = gr.Button("添加音频")
+            audio_status = gr.Textbox(label="音频状态", interactive=False)
+            audio_gallery = gr.Audio(label="已添加音频", type="filepath", interactive=False, visible=False)
+            register_btn = gr.Button("注册完成")
+            register_status = gr.Textbox(label="注册状态", interactive=False)
+            add_btn.click(
+                add_audio,
+                inputs=[text_input, audio_input, audio_list_state],
+                outputs=[audio_list_state, audio_status]
+            ).then(
+                update_gallery,
+                inputs=audio_list_state,
+                outputs=audio_gallery
+            )
+            register_btn.click(
+                register_keyword,
+                inputs=[text_input, audio_list_state],
+                outputs=register_status
+            )
+        with gr.Column(scale=2):
+            gr.Markdown("## 实时检测")
+            mic = gr.Audio(sources="microphone", streaming=True, label="实时监听")
+            state = gr.State(value=[])
+            audio_player = gr.Audio(label="唤醒提示", visible=False)
+            mic.stream(
+                streaming_detect_handler,
+                inputs=[mic, state, audio_player],
+                outputs=[state, audio_player],
+                time_limit=1000,
+                stream_every=0.05,
+            )
+if __name__ == "__main__":
+    demo.launch()