From 3d5e19792cd4bb510c2c0fc5749731d52b825c15 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期六, 08 六月 2024 18:43:35 +0800
Subject: [PATCH] fix bug

---
 funasr/models/llm_asr/model.py |   55 ++++++++++++++++++++++++++++++++++++++++++-------------
 1 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/funasr/models/llm_asr/model.py b/funasr/models/llm_asr/model.py
index 78d9340..697f78d 100644
--- a/funasr/models/llm_asr/model.py
+++ b/funasr/models/llm_asr/model.py
@@ -6,7 +6,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.cuda.amp import autocast
-
+import re
 from funasr.models.scama.utils import sequence_mask
 from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
 from funasr.models.ctc.ctc import CTC
@@ -532,7 +532,7 @@
         loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
         return loss, stats, weight
 
-    def data_template(self, data_in):
+    def data_template(self, data):
         system, user, assistant = [], [], []
         for i, item in enumerate(data):
             role = item["role"]
@@ -560,21 +560,31 @@
         user = contents["user"]
         assistant = contents["assistant"]
         pattern = re.compile(r"(<\|startofspeech\|>.*?<\|endofspeech\|>)")
-        input_ids, labels, fbank, fbank_lens, fbank_mask, fbank_beg = [], [], [], [], [], []
+        input_ids, labels, source_ids, target_ids, fbank, fbank_lens, fbank_mask, fbank_beg = (
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
 
         for i, (system_prompt, user_prompt, target_out) in enumerate(zip(system, user, assistant)):
 
             source_input = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
 
             splits = pattern.split(source_input)
-            source_ids = []
+            source_ids_i = []
             fbank_mask_i = []
             fbank_beg_i = []
             fbank_lens_i = []
+            # target_ids_i = []
             for k, sub_str in enumerate(splits):
                 if not sub_str.startswith("<|startofspeech|>"):
                     sub_token = tokenizer.encode(sub_str)
-                    source_ids += sub_token
+                    source_ids_i += sub_token
                     fbank_mask_i += [0] * len(sub_token)
                 else:
                     sub_str = sub_str.replace("<|startofspeech|>", "").replace(
@@ -600,14 +610,14 @@
                         olens = 1 + (olens - 3 + 2 * 1) // 2
                         sub_token_len = (olens - 1) // 2 + 1
                         sub_token = [0] * sub_token_len
-                        fbank_beg_i = [len(source_ids)]
-                        source_ids += sub_token
+                        fbank_beg_i = [len(source_ids_i)]
+                        source_ids_i += sub_token
                         fbank_mask_i += [1] * len(sub_token)
 
-            source_mask = [-100] * len(source_ids)
+            source_mask = [-100] * len(source_ids_i)
             target_out = f"{target_out}<|im_end|>"
             target_ids = tokenizer.encode(target_out)
-            input_ids += source_ids + target_ids
+            input_ids += source_ids_i + target_ids
             labels += source_mask + target_ids
             fbank_mask += fbank_mask_i
             fbank_beg.append(fbank_beg_i)
@@ -615,7 +625,7 @@
         input_ids = torch.tensor(input_ids, dtype=torch.int64)  # [: self.max_token_length]
         attention_mask = torch.tensor([1] * len(input_ids), dtype=torch.int32)
         labels = torch.tensor(labels, dtype=torch.int64)  # [: self.max_token_length]
-        source_ids = torch.tensor(source_ids, dtype=torch.int64)
+        source_ids = torch.tensor(source_ids_i, dtype=torch.int64)
         target_ids = torch.tensor(target_ids, dtype=torch.int64)
 
         fbank = speech[0, :, :]
@@ -653,7 +663,7 @@
         if kwargs.get("batch_size", 1) > 1:
             raise NotImplementedError("batch decoding is not implemented")
 
-        contents = self.data_template(data_in)
+        contents = self.data_template(data_in[0])
         output = self.data_load_speech(contents, tokenizer, frontend, **kwargs)
         batch = to_device(output, kwargs["device"])
 
@@ -667,7 +677,7 @@
 
         input_ids = batch["input_ids"]
         source_ids = batch["source_ids"]
-        if kwargs.get("tearchforing", False):
+        if not kwargs.get("tearchforing", False):
             input_ids = source_ids
         input_ids[input_ids < 0] = 0
         inputs_embeds = self.llm.model.get_input_embeddings()(input_ids)
@@ -695,6 +705,23 @@
                 generated_ids, skip_special_tokens=kwargs.get("skip_special_tokens", True)
             )[0]
             label = contents["assistant"][0]
+            loss = None
+        else:
+
+            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
+            )
+
+            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
 
         ibest_writer = None
         if kwargs.get("output_dir") is not None:
@@ -704,10 +731,12 @@
 
         results = []
         result_i = {"key": key[0], "text": response, "label": label}
+        if loss is not None:
+            result_i["loss"] = loss
         results.append(result_i)
 
         if ibest_writer is not None:
-            ibest_writer["text"][key[0]] = text
+            ibest_writer["text"][key[0]] = response
             ibest_writer["label"][key[0]] = label
 
         return results, meta_data

--
Gitblit v1.9.1