Dev gzf exp (#1654)

zhifu gao

2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e

 funasr/models/llm_asr_nar/model.py

@@ -12,6 +12,7 @@
from funasr.models.ctc.ctc import CTC
from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
from funasr.metrics.compute_acc import th_accuracy, compute_accuracy

# from funasr.models.e2e_asr_common import ErrorCalculator
from funasr.train_utils.device_funcs import force_gatherable
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
@@ -74,7 +75,11 @@
        hub = encoder_conf.get("hub", None)
        if hub == "funasr":
            from funasr import AutoModel
            init_param_path = encoder_conf.get("init_param_path", "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch")

            init_param_path = encoder_conf.get(
                "init_param_path",
                "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
            )
            model = AutoModel(model=init_param_path, model_revision="master")
            # frontend = model.kwargs.get("frontend")
            model.model.decoder = None
@@ -115,7 +120,6 @@
        
        self.adaptor = adaptor
        
        
        self.blank_id = blank_id
        self.sos = sos if sos is not None else vocab_size - 1
        self.eos = eos if eos is not None else vocab_size - 1
@@ -124,7 +128,6 @@
        self.specaug = specaug
        self.normalize = normalize
        self.encoder = encoder


        self.criterion_att = LabelSmoothingLoss(
            size=vocab_size,
@@ -192,12 +195,15 @@
                batch_size, token_num, dims = inputs_embeds.shape
                _, l, _ = encoder_out.shape
                encoder_outs_pad = F.pad(encoder_out, (0, 0, token_num-l-1, 1, 0, 0), value=0.0)
                inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (1.0-audio_mask[:, :, None])
                inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (
                    1.0 - audio_mask[:, :, None]
                )
                inputs_embeds = F.pad(inputs_embeds[:, 1:, :], (0, 0, 0, 1, 0, 0), value=0.0)

        model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids)
        model_outputs = self.llm(
            inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids
        )
        loss = model_outputs.loss


        stats = {}
        with torch.no_grad():
@@ -214,7 +220,10 @@
        return loss, stats, weight
    
    def encode(
        self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs,
        self,
        speech: torch.Tensor,
        speech_lengths: torch.Tensor,
        **kwargs,
    ):
    
        audio_mask = kwargs.get("audio_mask", None)
@@ -227,15 +236,16 @@
        enc, enc_lens = self.audio_encoder.encode(**batch)
        with autocast(False):
            enc_mask = sequence_mask(enc_lens, enc.size(1), device=enc.device)[:, None, :]
            pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor(enc,
            pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor(
                enc,
                                                                               mask=enc_mask,
                                                                               target_label_length=audio_token_lengths,
                                                                               )

        return pre_acoustic_embeds, pre_token_length


    def inference(self,
    def inference(
        self,
                  data_in,
                  data_lengths=None,
                  key: list = None,
@@ -249,10 +259,10 @@
        if kwargs.get("batch_size", 1) > 1:
            raise NotImplementedError("batch decoding is not implemented")


        
        meta_data = {}
        if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank":  # fbank
        if (
            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
        ):  # fbank
            speech, speech_lengths = data_in, data_lengths
            if len(speech.shape) < 3:
                speech = speech[None, :, :]
@@ -261,9 +271,13 @@
        else:
            # extract fbank feats
            time1 = time.perf_counter()
            audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000),
            audio_sample_list = load_audio_text_image_video(
                data_in,
                fs=frontend.fs,
                audio_fs=kwargs.get("fs", 16000),
                                                            data_type=kwargs.get("data_type", "sound"),
                                                            tokenizer=None)
                tokenizer=None,
            )
            if len(kwargs.get("data_type", [])) > 1:
                audio_sample_list, text_token_int_list = audio_sample_list
                text_token_int = text_token_int_list[0].replace(" ", "")
@@ -272,27 +286,30 @@
                text_token_int = None
            time2 = time.perf_counter()
            meta_data["load_data"] = f"{time2 - time1:0.3f}"
            speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"),
                                                   frontend=frontend)
            speech, speech_lengths = extract_fbank(
                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
            )
            time3 = time.perf_counter()
            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
            meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
            meta_data["batch_data_time"] = (
                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
            )
        
        speech = speech.to(device=kwargs["device"])
        speech_lengths = speech_lengths.to(device=kwargs["device"])
        
        # Encoder
        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, text_token_int=text_token_int)
        encoder_out, encoder_out_lens = self.encode(
            speech, speech_lengths, text_token_int=text_token_int
        )

        # adaptor
        encoder_out = self.adaptor(encoder_out)
        
    
        prompt_pre = "USER: \nINSTRUCTION: {}\nINPUT: ".format(prompt)
        prompt_ids = tokenizer.encode(prompt_pre)
        prompt_length = len(prompt_ids)
        prompt_ids = torch.tensor(prompt_ids, dtype=torch.int64).to(kwargs["device"])


        if hasattr(self.llm.model, "embed_tokens"):
            inputs_embeds = self.llm.model.embed_tokens(prompt_ids)
@@ -301,8 +318,12 @@
        else:
            inputs_embeds = self.llm.model.model.model.embed_tokens(prompt_ids)

        inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out), dim=1)  # [prompt, audio]
        attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(kwargs["device"])
        inputs_embeds = torch.cat(
            (inputs_embeds[None, :, :], encoder_out), dim=1
        )  # [prompt, audio]
        attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(
            kwargs["device"]
        )
        
        # model_outputs = self.llm.generate(
        #     inputs_embeds=inputs_embeds,
@@ -321,12 +342,13 @@
        #     pad_token_id=tokenizer.pad_token_id
        # )


        model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None)
        model_outputs = self.llm(
            inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None
        )
        preds = torch.argmax(model_outputs.logits, -1)
        text = tokenizer.batch_decode(preds, add_special_tokens=False, skip_special_tokens=True)

        text = text[0].split(': ')[-1]
        text = text[0].split(": ")[-1]
        text = text.strip()
        
        # preds = torch.argmax(model_outputs.logits, -1)
@@ -343,9 +365,6 @@

        if ibest_writer is not None:
            ibest_writer["text"][key[0]] = text
        
        
        
        
        return results, meta_data

@@ -404,8 +423,11 @@
        hub = encoder_conf.get("hub", None)
        if hub == "funasr":
            from funasr import AutoModel
            init_param_path = encoder_conf.get("init_param_path",
                                               "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch")

            init_param_path = encoder_conf.get(
                "init_param_path",
                "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
            )
            model = AutoModel(model=init_param_path, model_revision="master")
            # frontend = model.kwargs.get("frontend")
            model.model.decoder = None
@@ -477,9 +499,7 @@
            if ctc_conf is None:
                ctc_conf = {}
    
            ctc = CTC(
                odim=vocab_size, encoder_output_size=adaptor_conf["encoder_dim"], **ctc_conf
            )
            ctc = CTC(odim=vocab_size, encoder_output_size=adaptor_conf["encoder_dim"], **ctc_conf)
        self.ctc_weight = ctc_weight
        self.ctc = ctc
    
@@ -518,13 +538,10 @@
        enc, enc_lens = outs[0], outs[1]
        encoder_out, encoder_out_lens, loss_pre = outs[2], outs[3], outs[4]
        

        # decoder: CTC branch
        
        if self.ctc_weight != 0.0:
            loss_ctc, cer_ctc = self._calc_ctc_loss(
                enc, enc_lens, text, text_lengths
            )
            loss_ctc, cer_ctc = self._calc_ctc_loss(enc, enc_lens, text, text_lengths)
    
            # Collect CTC branch stats
            stats["loss_ctc"] = torch.clone(loss_ctc.detach()) if loss_ctc is not None else None
@@ -549,15 +566,25 @@
                prompt_bos_length = prompt_bos_length[0].item()
                batch_size, token_num, dims = inputs_embeds.shape
                _, l, _ = encoder_out.shape
                encoder_outs_pad = F.pad(encoder_out, (0, 0, prompt_bos_length, token_num - prompt_bos_length - l, 0, 0), value=0.0)
                inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (1.0 - audio_mask[:, :, None])
                inputs_embeds = F.pad(inputs_embeds[:, 1:, :], (0, 0, 0, 1, 0, 0), value=0.0) # [prompt, input, pad, target, 0.0]
                encoder_outs_pad = F.pad(
                    encoder_out,
                    (0, 0, prompt_bos_length, token_num - prompt_bos_length - l, 0, 0),
                    value=0.0,
                )
                inputs_embeds = encoder_outs_pad * audio_mask[:, :, None] + inputs_embeds * (
                    1.0 - audio_mask[:, :, None]
                )
                inputs_embeds = F.pad(
                    inputs_embeds[:, 1:, :], (0, 0, 0, 1, 0, 0), value=0.0
                )  # [prompt, input, pad, target, 0.0]
        
        # labels_ids: [bos, prompt, input, target, eos] -> [-1, -1, input, target, eos]
        # loss:
        # inputs_embeds[:-1] -> [prompt, input, pad, target]
        # labels_ids[1:] ->  [prompt, input, target, eos] -> [-1, input, target, eos];
        model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids)
        model_outputs = self.llm(
            inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids
        )
        loss_llm = model_outputs.loss
        stats["loss_llm"] = torch.clone(loss_llm.detach())
        if self.ctc_weight > 0.0:
@@ -580,7 +607,10 @@
        return loss, stats, weight
    
    def encode(
        self, speech: torch.Tensor, speech_lengths: torch.Tensor, **kwargs,
        self,
        speech: torch.Tensor,
        speech_lengths: torch.Tensor,
        **kwargs,
    ):
        
        audio_mask = kwargs.get("audio_mask", None)
@@ -593,13 +623,16 @@
        enc, enc_lens = self.audio_encoder.encode(**batch)
        with autocast(False):
            enc_mask = sequence_mask(enc_lens, enc.size(1), device=enc.device)[:, None, :]
            pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor(enc,
            pre_acoustic_embeds, pre_token_length, _, _ = self.audio_encoder.predictor(
                enc,
                                                                                       mask=enc_mask,
                                                                                       target_label_length=audio_token_lengths,
                                                                                       )
            loss_pre = 0.0
            if audio_token_lengths is not None:
                loss_pre = self.criterion_pre(audio_token_lengths.type_as(pre_token_length), pre_token_length)
                loss_pre = self.criterion_pre(
                    audio_token_lengths.type_as(pre_token_length), pre_token_length
                )
        
        return enc, enc_lens, pre_acoustic_embeds, pre_token_length, loss_pre

@@ -620,7 +653,8 @@
            cer_ctc = self.error_calculator(ys_hat.cpu(), ys_pad.cpu(), is_ctc=True)
        return loss_ctc, cer_ctc
    
    def inference(self,
    def inference(
        self,
                  data_in,
                  data_lengths=None,
                  key: list = None,
@@ -635,7 +669,9 @@
            raise NotImplementedError("batch decoding is not implemented")
        
        meta_data = {}
        if isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank":  # fbank
        if (
            isinstance(data_in, torch.Tensor) and kwargs.get("data_type", "sound") == "fbank"
        ):  # fbank
            speech, speech_lengths = data_in, data_lengths
            if len(speech.shape) < 3:
                speech = speech[None, :, :]
@@ -644,9 +680,13 @@
        else:
            # extract fbank feats
            time1 = time.perf_counter()
            audio_sample_list = load_audio_text_image_video(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000),
            audio_sample_list = load_audio_text_image_video(
                data_in,
                fs=frontend.fs,
                audio_fs=kwargs.get("fs", 16000),
                                                            data_type=kwargs.get("data_type", "sound"),
                                                            tokenizer=None)
                tokenizer=None,
            )
            if len(kwargs.get("data_type", [])) > 1:
                audio_sample_list, text_token_int_list = audio_sample_list
                text_token_int = text_token_int_list[0]
@@ -657,11 +697,14 @@
                text_token_int = None
            time2 = time.perf_counter()
            meta_data["load_data"] = f"{time2 - time1:0.3f}"
            speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"),
                                                   frontend=frontend)
            speech, speech_lengths = extract_fbank(
                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
            )
            time3 = time.perf_counter()
            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
            meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
            meta_data["batch_data_time"] = (
                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
            )
        
        speech = speech.to(device=kwargs["device"])
        speech_lengths = speech_lengths.to(device=kwargs["device"])
@@ -691,8 +734,12 @@
            inputs_embeds = self.llm.model.model.model.embed_tokens(prompt_ids)
        
        # inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out, pad[None, :, :]), dim=1)  # [prompt, audio, pad]
        inputs_embeds = torch.cat((inputs_embeds[None, :, :], encoder_out), dim=1)  # [prompt, audio]
        attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(kwargs["device"])
        inputs_embeds = torch.cat(
            (inputs_embeds[None, :, :], encoder_out), dim=1
        )  # [prompt, audio]
        attention_mask = torch.ones(inputs_embeds.size()[:-1], dtype=torch.long).to(
            kwargs["device"]
        )
        
        # model_outputs = self.llm.generate(
        #     inputs_embeds=inputs_embeds,
@@ -711,11 +758,13 @@
        #     pad_token_id=tokenizer.pad_token_id
        # )
        
        model_outputs = self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None)
        model_outputs = self.llm(
            inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=None
        )
        preds = torch.argmax(model_outputs.logits, -1)
        text = tokenizer.batch_decode(preds, add_special_tokens=False, skip_special_tokens=True)
        
        text = text[0].split(':')[-1]
        text = text[0].split(":")[-1]
        text = text.strip()
        if text.startswith("Please\n "):
            text = text.replace("Please\n ", "")