From 3df109adfccedeb134dea4ba2ea9a2da89872048 Mon Sep 17 00:00:00 2001
From: Isuxiz Slidder <48672727+Isuxiz@users.noreply.github.com>
Date: 星期一, 31 三月 2025 17:51:52 +0800
Subject: [PATCH] Update model.py to fix "IndexError: index 1 is out of bounds for dimension 1 with size 0" (#2454)

---
 funasr/models/sense_voice/model.py |   80 ++++++++++++++++++++++++++++++++++++++-
 1 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/funasr/models/sense_voice/model.py b/funasr/models/sense_voice/model.py
index 25e9faf..c5d0e59 100644
--- a/funasr/models/sense_voice/model.py
+++ b/funasr/models/sense_voice/model.py
@@ -19,6 +19,7 @@
 
 
 from funasr.models.paraformer.search import Hypothesis
+from .utils.ctc_alignment import ctc_forced_align
 
 
 class SinusoidalPositionEncoder(torch.nn.Module):
@@ -555,7 +556,8 @@
         ilens: torch.Tensor,
     ):
         """Embed positions in tensor."""
-        masks = sequence_mask(ilens, device=ilens.device)[:, None, :]
+        maxlen = xs_pad.shape[1]
+        masks = sequence_mask(ilens, maxlen=maxlen, device=ilens.device)[:, None, :]
 
         xs_pad *= self.output_size() ** 0.5
 
@@ -856,6 +858,8 @@
 
         use_itn = kwargs.get("use_itn", False)
         textnorm = kwargs.get("text_norm", None)
+        output_timestamp = kwargs.get("output_timestamp", False)
+
         if textnorm is None:
             textnorm = "withitn" if use_itn else "woitn"
         textnorm_query = self.embed(
@@ -904,13 +908,81 @@
             # Change integer-ids to tokens
             text = tokenizer.decode(token_int)
 
-            result_i = {"key": key[i], "text": text}
-            results.append(result_i)
+            # result_i = {"key": key[i], "text": text}
+            # results.append(result_i)
 
             if ibest_writer is not None:
                 ibest_writer["text"][key[i]] = text
 
+            if output_timestamp:
+                from itertools import groupby
+
+                timestamp = []
+                tokens = tokenizer.text2tokens(text)[4:]
+                token_back_to_id = tokenizer.tokens2ids(tokens)
+                token_ids = []
+                for tok_ls in token_back_to_id:
+                    if tok_ls: token_ids.extend(tok_ls)
+                    else: token_ids.append(124)
+
+                if len(token_ids) == 0:
+                    result_i = {"key": key[i], "text": text}
+                    results.append(result_i)
+                    continue
+
+                logits_speech = self.ctc.softmax(encoder_out)[i, 4 : encoder_out_lens[i].item(), :]
+                pred = logits_speech.argmax(-1).cpu()
+                logits_speech[pred == self.blank_id, self.blank_id] = 0
+                align = ctc_forced_align(
+                    logits_speech.unsqueeze(0).float(),
+                    torch.Tensor(token_ids).unsqueeze(0).long().to(logits_speech.device),
+                    (encoder_out_lens[i] - 4).long(),
+                    torch.tensor(len(token_ids)).unsqueeze(0).long().to(logits_speech.device),
+                    ignore_id=self.ignore_id,
+                )
+                pred = groupby(align[0, : encoder_out_lens[i]])
+                _start = 0
+                token_id = 0
+                ts_max = encoder_out_lens[i] - 4
+                for pred_token, pred_frame in pred:
+                    _end = _start + len(list(pred_frame))
+                    if pred_token != 0:
+                        ts_left = max((_start * 60 - 30) / 1000, 0)
+                        ts_right = min((_end * 60 - 30) / 1000, (ts_max * 60 - 30) / 1000)
+                        timestamp.append([tokens[token_id], ts_left, ts_right])
+                        token_id += 1
+                    _start = _end
+                timestamp = self.post(timestamp)
+                result_i = {"key": key[i], "text": text, "timestamp": timestamp}
+                results.append(result_i)
+            else:
+                result_i = {"key": key[i], "text": text}
+                results.append(result_i)
         return results, meta_data
+
+    def post(self, timestamp):
+        timestamp_new = []
+        prev_word = None
+        for i, t in enumerate(timestamp):
+            word, start, end = t
+            start = int(start * 1000)
+            end = int(end * 1000)
+            if word == "鈻�":
+                continue
+            if i == 0:
+                # timestamp_new.append([word, start, end])
+                timestamp_new.append([start, end])
+            elif word.startswith("鈻�"):
+                word = word[1:]
+                timestamp_new.append([start, end])
+            elif prev_word is not None and prev_word.isalpha() and prev_word.isascii() and word.isalpha() and word.isascii():
+                prev_word += word
+                timestamp_new[-1][1] = end
+            else:
+                # timestamp_new[-1][0] += word
+                timestamp_new.append([start, end])
+            prev_word = word
+        return timestamp_new
 
     def export(self, **kwargs):
         from .export_meta import export_rebuild_model
@@ -919,3 +991,5 @@
             kwargs["max_seq_len"] = 512
         models = export_rebuild_model(model=self, **kwargs)
         return models
+
+        return results, meta_data

--
Gitblit v1.9.1