Dev gzf exp (#1654)

zhifu gao

2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e

 funasr/models/paraformer/cif_predictor.py

@@ -12,9 +12,20 @@
from funasr.models.transformer.utils.nets_utils import make_pad_mask

from torch.cuda.amp import autocast





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

class CifPredictor(torch.nn.Module):

    def __init__(self, idim, l_order, r_order, threshold=1.0, dropout=0.1, smooth_factor=1.0, noise_threshold=0, tail_threshold=0.45):

    def __init__(

        self,

        idim,

        l_order,

        r_order,

        threshold=1.0,

        dropout=0.1,

        smooth_factor=1.0,

        noise_threshold=0,

        tail_threshold=0.45,

    ):

        super().__init__()



        self.pad = torch.nn.ConstantPad1d((l_order, r_order), 0)

@@ -26,8 +37,15 @@
        self.noise_threshold = noise_threshold

        self.tail_threshold = tail_threshold



    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,

    ):

    

        with autocast(False):

            h = hidden

@@ -58,7 +76,9 @@
            if target_length is not None:

                alphas *= (target_length / token_num)[:, None].repeat(1, alphas.size(1))

            elif self.tail_threshold > 0.0:

                hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=mask)

                hidden, alphas, token_num = self.tail_process_fn(

                    hidden, alphas, token_num, mask=mask

                )

                

            acoustic_embeds, cif_peak = cif(hidden, alphas, self.threshold)

            

@@ -91,10 +111,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



@@ -117,11 +136,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)

@@ -132,17 +155,24 @@
        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", "CifPredictorV2")

class CifPredictorV2(torch.nn.Module):

    def __init__(self,

    def __init__(

        self,

                 idim,

                 l_order,

                 r_order,

@@ -169,8 +199,15 @@
        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf

        self.tail_mask = tail_mask



    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,

    ):

        

        with autocast(False):

            h = hidden

@@ -200,9 +237,13 @@
                alphas *= (target_length / token_num)[:, None].repeat(1, alphas.size(1))

            elif self.tail_threshold > 0.0:

                if self.tail_mask:

                    hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=mask)

                    hidden, alphas, token_num = self.tail_process_fn(

                        hidden, alphas, token_num, mask=mask

                    )

                else:

                    hidden, alphas, token_num = self.tail_process_fn(hidden, alphas, token_num, mask=None)

                    hidden, alphas, token_num = self.tail_process_fn(

                        hidden, alphas, token_num, mask=None

                    )

    

            acoustic_embeds, cif_peak = cif(hidden, alphas, self.threshold)

            if target_length is None and self.tail_threshold > 0.0:

@@ -287,7 +328,9 @@
             return hidden, torch.stack(token_length, 0), None, None

        list_ls = []

        for b in range(batch_size):

            pad_frames = torch.zeros((max_token_len - token_length[b], hidden_size), device=alphas.device)

            pad_frames = torch.zeros(

                (max_token_len - token_length[b], hidden_size), device=alphas.device

            )

            if token_length[b] == 0:

                list_ls.append(pad_frames)

            else:

@@ -299,7 +342,6 @@
        cache["cif_hidden"] = torch.stack(cache_hiddens, axis=0)

        cache["cif_hidden"] = torch.unsqueeze(cache["cif_hidden"], axis=0)

        return torch.stack(list_ls, 0), torch.stack(token_length, 0), None, None





    def tail_process_fn(self, hidden, alphas, token_num=None, mask=None):

        b, t, d = hidden.size()

@@ -327,9 +369,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



@@ -352,11 +394,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)

@@ -367,13 +413,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", "CifPredictorV2Export")

class CifPredictorV2Export(torch.nn.Module):

@@ -388,7 +440,9 @@
        self.noise_threshold = model.noise_threshold

        self.tail_threshold = model.tail_threshold

    

    def forward(self, hidden: torch.Tensor,

    def forward(

        self,

        hidden: torch.Tensor,

                mask: torch.Tensor,

                ):

        alphas, token_num = self.forward_cnn(hidden, mask)

@@ -399,7 +453,9 @@
        

        return acoustic_embeds, token_num, alphas, cif_peak

    

    def forward_cnn(self, hidden: torch.Tensor,

    def forward_cnn(

        self,

        hidden: torch.Tensor,

                    mask: torch.Tensor,

                    ):

        h = hidden

@@ -439,6 +495,7 @@
        

        return hidden, alphas, token_num_floor





@torch.jit.script

def cif_export(hidden, alphas, threshold: float):

    batch_size, len_time, hidden_size = hidden.size()

@@ -453,25 +510,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)

@@ -495,7 +554,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)

@@ -524,19 +583,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)

@@ -566,10 +623,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