decoding

游雁

2024-06-11 d43f77408b8f3e169c59dfb6b6d82e45e6b91714

 funasr/models/llm_asr/model.py

@@ -684,6 +684,11 @@
        # audio encoder
        speech = batch["speech"]
        speech_lengths = batch["speech_lengths"][:, 0]
        # fp16
        if kwargs.get("fp16", False):
            speech = speech.to(torch.float16)
        elif kwargs.get("bf16", False):
            speech = speech.to(torch.bfloat16)
        encoder_out, encoder_out_lens = self.audio_encoder(speech.permute(0, 2, 1), speech_lengths)

        # audio_adaptor
@@ -706,37 +711,50 @@
                batch_idx, :min_len, :
            ]

        label = contents["assistant"][0]
        if not kwargs.get("tearchforing", False):
        llm_dtype = kwargs.get("llm_dtype", "fp32")
        if llm_dtype == "fp32":
            llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype
            llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype

            generated_ids = self.llm.generate(
                inputs_embeds=inputs_embeds, max_new_tokens=kwargs.get("max_length", 512)
            )
            # generated_ids = [
            #     output_ids[len(input_id) :]
            #     for input_id, output_ids in zip(input_ids, generated_ids)
            # ]
            response = tokenizer.batch_decode(
                generated_ids, skip_special_tokens=kwargs.get("skip_special_tokens", True)
            )[0]
        dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
        with torch.cuda.amp.autocast(
            enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype]
        ):
            label = contents["assistant"][0]
            self.llm = self.llm.to(dtype_map[llm_dtype])
            inputs_embeds = inputs_embeds.to(dtype_map[llm_dtype])

            loss = None
        else:
            if not kwargs.get("tearchforing", False):

            labels_ids = batch["labels_ids"]
            labels_ids[labels_ids == -1] = -100
            attention_mask = batch.get("attention_mask", None)
            model_outputs = self.llm(
                inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids
            )
                generated_ids = self.llm.generate(
                    inputs_embeds=inputs_embeds, max_new_tokens=kwargs.get("max_length", 512)
                )
                # generated_ids = [
                #     output_ids[len(input_id) :]
                #     for input_id, output_ids in zip(input_ids, generated_ids)
                # ]
                response = tokenizer.batch_decode(
                    generated_ids, skip_special_tokens=kwargs.get("skip_special_tokens", True)
                )[0]

            preds = torch.argmax(model_outputs.logits, -1)[:, source_ids.shape[1] :]
            response = tokenizer.batch_decode(
                preds,
                add_special_tokens=False,
                skip_special_tokens=kwargs.get("skip_special_tokens", True),
            )[0]
            loss = model_outputs.loss.item()
                loss = None
            else:

                labels_ids = batch["labels_ids"]
                labels_ids[labels_ids == -1] = -100
                attention_mask = batch.get("attention_mask", None)
                # attention_mask = attention_mask.to(dtype_map[llm_dtype])
                model_outputs = self.llm(
                    inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels_ids
                )

                preds = torch.argmax(model_outputs.logits, -1)[:, source_ids.shape[1] :]
                response = tokenizer.batch_decode(
                    preds,
                    add_special_tokens=False,
                    skip_special_tokens=kwargs.get("skip_special_tokens", True),
                )[0]
                loss = model_outputs.loss.item()

        ibest_writer = None
        if kwargs.get("output_dir") is not None: