Dev gzf exp (#1654)

zhifu gao

2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e

 funasr/models/bicif_paraformer/cif_predictor.py

@@ -14,7 +14,7 @@
    def __init__(self, normalize_length=False):

        super(mae_loss, self).__init__()

        self.normalize_length = normalize_length

        self.criterion = torch.nn.L1Loss(reduction='sum')

        self.criterion = torch.nn.L1Loss(reduction="sum")



    def forward(self, token_length, pre_token_length):

        loss_token_normalizer = token_length.size(0)

@@ -43,19 +43,17 @@
        list_fires.append(integrate)



        fire_place = integrate >= threshold

        integrate = torch.where(fire_place,

                                integrate - torch.ones([batch_size], device=hidden.device),

                                integrate)

        cur = torch.where(fire_place,

                          distribution_completion,

                          alpha)

        integrate = torch.where(

            fire_place, integrate - torch.ones([batch_size], device=hidden.device), integrate

        )

        cur = torch.where(fire_place, distribution_completion, alpha)

        remainds = alpha - cur



        frame += cur[:, None] * hidden[:, t, :]

        list_frames.append(frame)

        frame = torch.where(fire_place[:, None].repeat(1, hidden_size),

                            remainds[:, None] * hidden[:, t, :],

                            frame)

        frame = torch.where(

            fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame

        )



    fires = torch.stack(list_fires, 1)

    frames = torch.stack(list_frames, 1)

@@ -85,16 +83,20 @@
        list_fires.append(integrate)



        fire_place = integrate >= threshold

        integrate = torch.where(fire_place,

        integrate = torch.where(

            fire_place,

                                integrate - torch.ones([batch_size], device=alphas.device)*threshold,

                                integrate)

            integrate,

        )



    fires = torch.stack(list_fires, 1)

    return fires





@tables.register("predictor_classes", "CifPredictorV3")

class CifPredictorV3(torch.nn.Module):

    def __init__(self,

    def __init__(

        self,

                 idim,

                 l_order,

                 r_order,

@@ -128,18 +130,27 @@
        self.upsample_times = upsample_times

        self.upsample_type = upsample_type

        self.use_cif1_cnn = use_cif1_cnn

        if self.upsample_type == 'cnn':

            self.upsample_cnn = torch.nn.ConvTranspose1d(idim, idim, self.upsample_times, self.upsample_times)

        if self.upsample_type == "cnn":

            self.upsample_cnn = torch.nn.ConvTranspose1d(

                idim, idim, self.upsample_times, self.upsample_times

            )

            self.cif_output2 = torch.nn.Linear(idim, 1)

        elif self.upsample_type == 'cnn_blstm':

            self.upsample_cnn = torch.nn.ConvTranspose1d(idim, idim, self.upsample_times, self.upsample_times)

            self.blstm = torch.nn.LSTM(idim, idim, 1, bias=True, batch_first=True, dropout=0.0, bidirectional=True)

        elif self.upsample_type == "cnn_blstm":

            self.upsample_cnn = torch.nn.ConvTranspose1d(

                idim, idim, self.upsample_times, self.upsample_times

            )

            self.blstm = torch.nn.LSTM(

                idim, idim, 1, bias=True, batch_first=True, dropout=0.0, bidirectional=True

            )

            self.cif_output2 = torch.nn.Linear(idim*2, 1)

        elif self.upsample_type == 'cnn_attn':

            self.upsample_cnn = torch.nn.ConvTranspose1d(idim, idim, self.upsample_times, self.upsample_times)

        elif self.upsample_type == "cnn_attn":

            self.upsample_cnn = torch.nn.ConvTranspose1d(

                idim, idim, self.upsample_times, self.upsample_times

            )

            from funasr.models.transformer.encoder import EncoderLayer as TransformerEncoderLayer

            from funasr.models.transformer.attention import MultiHeadedAttention

            from funasr.models.transformer.positionwise_feed_forward import PositionwiseFeedForward



            positionwise_layer_args = (

                idim,

                idim*2,

@@ -147,9 +158,7 @@
            )

            self.self_attn = TransformerEncoderLayer(

                idim,

                MultiHeadedAttention(

                    4, idim, 0.1

                ),

                MultiHeadedAttention(4, idim, 0.1),

                PositionwiseFeedForward(*positionwise_layer_args),

                0.1,

                True, #normalize_before,

@@ -159,8 +168,15 @@
        self.smooth_factor2 = smooth_factor2

        self.noise_threshold2 = noise_threshold2



    def forward(self, hidden, target_label=None, mask=None, ignore_id=-1, mask_chunk_predictor=None,

                target_label_length=None):

    def forward(

        self,

        hidden,

        target_label=None,

        mask=None,

        ignore_id=-1,

        mask_chunk_predictor=None,

        target_label_length=None,

    ):

        h = hidden

        context = h.transpose(1, 2)

        queries = self.pad(context)

@@ -171,14 +187,14 @@
            _output = context

        else:

            _output = output

        if self.upsample_type == 'cnn':

        if self.upsample_type == "cnn":

            output2 = self.upsample_cnn(_output)

            output2 = output2.transpose(1,2)

        elif self.upsample_type == 'cnn_blstm':

        elif self.upsample_type == "cnn_blstm":

            output2 = self.upsample_cnn(_output)

            output2 = output2.transpose(1,2)

            output2, (_, _) = self.blstm(output2)

        elif self.upsample_type == 'cnn_attn':

        elif self.upsample_type == "cnn_attn":

            output2 = self.upsample_cnn(_output)

            output2 = output2.transpose(1,2)

            output2, _ = self.self_attn(output2, mask)

@@ -187,7 +203,11 @@
        alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2)

        # repeat the mask in T demension to match the upsampled length

        if mask is not None:

            mask2 = mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1)

            mask2 = (

                mask.repeat(1, self.upsample_times, 1)

                .transpose(-1, -2)

                .reshape(alphas2.shape[0], -1)

            )

            mask2 = mask2.unsqueeze(-1)

            alphas2 = alphas2 * mask2

        alphas2 = alphas2.squeeze(-1)

@@ -236,14 +256,14 @@
            _output = context

        else:

            _output = output

        if self.upsample_type == 'cnn':

        if self.upsample_type == "cnn":

            output2 = self.upsample_cnn(_output)

            output2 = output2.transpose(1,2)

        elif self.upsample_type == 'cnn_blstm':

        elif self.upsample_type == "cnn_blstm":

            output2 = self.upsample_cnn(_output)

            output2 = output2.transpose(1,2)

            output2, (_, _) = self.blstm(output2)

        elif self.upsample_type == 'cnn_attn':

        elif self.upsample_type == "cnn_attn":

            output2 = self.upsample_cnn(_output)

            output2 = output2.transpose(1,2)

            output2, _ = self.self_attn(output2, mask)

@@ -251,7 +271,11 @@
        alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2)

        # repeat the mask in T demension to match the upsampled length

        if mask is not None:

            mask2 = mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1)

            mask2 = (

                mask.repeat(1, self.upsample_times, 1)

                .transpose(-1, -2)

                .reshape(alphas2.shape[0], -1)

            )

            mask2 = mask2.unsqueeze(-1)

            alphas2 = alphas2 * mask2

        alphas2 = alphas2.squeeze(-1)

@@ -289,9 +313,9 @@


        return hidden, alphas, token_num_floor



    def gen_frame_alignments(self,

                             alphas: torch.Tensor = None,

                             encoder_sequence_length: torch.Tensor = None):

    def gen_frame_alignments(

        self, alphas: torch.Tensor = None, encoder_sequence_length: torch.Tensor = None

    ):

        batch_size, maximum_length = alphas.size()

        int_type = torch.int32



@@ -314,11 +338,15 @@
        index_div = torch.floor(torch.true_divide(alphas_cumsum, index)).type(int_type)

        index_div_bool_zeros = index_div.eq(0)

        index_div_bool_zeros_count = torch.sum(index_div_bool_zeros, dim=-1) + 1

        index_div_bool_zeros_count = torch.clamp(index_div_bool_zeros_count, 0, encoder_sequence_length.max())

        index_div_bool_zeros_count = torch.clamp(

            index_div_bool_zeros_count, 0, encoder_sequence_length.max()

        )

        token_num_mask = (~make_pad_mask(token_num, maxlen=max_token_num)).to(token_num.device)

        index_div_bool_zeros_count *= token_num_mask



        index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat(1, 1, maximum_length)

        index_div_bool_zeros_count_tile = index_div_bool_zeros_count[:, :, None].repeat(

            1, 1, maximum_length

        )

        ones = torch.ones_like(index_div_bool_zeros_count_tile)

        zeros = torch.zeros_like(index_div_bool_zeros_count_tile)

        ones = torch.cumsum(ones, dim=2)

@@ -329,13 +357,19 @@
        index_div_bool_zeros_count_tile = 1 - index_div_bool_zeros_count_tile_bool.type(int_type)

        index_div_bool_zeros_count_tile_out = torch.sum(index_div_bool_zeros_count_tile, dim=1)

        index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out.type(int_type)

        predictor_mask = (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max())).type(

            int_type).to(encoder_sequence_length.device)

        predictor_mask = (

            (~make_pad_mask(encoder_sequence_length, maxlen=encoder_sequence_length.max()))

            .type(int_type)

            .to(encoder_sequence_length.device)

        )

        index_div_bool_zeros_count_tile_out = index_div_bool_zeros_count_tile_out * predictor_mask



        predictor_alignments = index_div_bool_zeros_count_tile_out

        predictor_alignments_length = predictor_alignments.sum(-1).type(encoder_sequence_length.dtype)

        predictor_alignments_length = predictor_alignments.sum(-1).type(

            encoder_sequence_length.dtype

        )

        return predictor_alignments.detach(), predictor_alignments_length.detach()





@tables.register("predictor_classes", "CifPredictorV3Export")

class CifPredictorV3Export(torch.nn.Module):

@@ -357,7 +391,9 @@
        self.smooth_factor2 = model.smooth_factor2

        self.noise_threshold2 = model.noise_threshold2

    

    def forward(self, hidden: torch.Tensor,

    def forward(

        self,

        hidden: torch.Tensor,

                mask: torch.Tensor,

                ):

        h = hidden

@@ -393,7 +429,9 @@
        alphas2 = torch.sigmoid(self.cif_output2(output2))

        alphas2 = torch.nn.functional.relu(alphas2 * self.smooth_factor2 - self.noise_threshold2)

        

        mask = mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1)

        mask = (

            mask.repeat(1, self.upsample_times, 1).transpose(-1, -2).reshape(alphas2.shape[0], -1)

        )

        mask = mask.unsqueeze(-1)

        alphas2 = alphas2 * mask

        alphas2 = alphas2.squeeze(-1)

@@ -440,25 +478,27 @@
    

    for t in range(len_time):

        alpha = alphas[:, t]

        distribution_completion = torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate

        distribution_completion = (

            torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device) - integrate

        )

        

        integrate += alpha

        list_fires.append(integrate)

        

        fire_place = integrate >= threshold

        integrate = torch.where(fire_place,

        integrate = torch.where(

            fire_place,

                                integrate - torch.ones([batch_size], dtype=alphas.dtype, device=hidden.device),

                                integrate)

        cur = torch.where(fire_place,

                          distribution_completion,

                          alpha)

            integrate,

        )

        cur = torch.where(fire_place, distribution_completion, alpha)

        remainds = alpha - cur

        

        frame += cur[:, None] * hidden[:, t, :]

        list_frames.append(frame)

        frame = torch.where(fire_place[:, None].repeat(1, hidden_size),

                            remainds[:, None] * hidden[:, t, :],

                            frame)

        frame = torch.where(

            fire_place[:, None].repeat(1, hidden_size), remainds[:, None] * hidden[:, t, :], frame

        )

    

    fires = torch.stack(list_fires, 1)

    frames = torch.stack(list_frames, 1)

@@ -493,9 +533,11 @@
        list_fires.append(integrate)

        

        fire_place = integrate >= threshold

        integrate = torch.where(fire_place,

        integrate = torch.where(

            fire_place,

                                integrate - torch.ones([batch_size], device=alphas.device) * threshold,

                                integrate)

            integrate,

        )

    

    fires = torch.stack(list_fires, 1)

    return fires