From 001a66bbfe8093d2ff7336eeebdb0198b498dce9 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 11 二月 2025 10:08:19 +0800
Subject: [PATCH] oom fix

---
 funasr/models/sense_voice/model.py |   31 ++++++++++++++++---------------
 1 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/funasr/models/sense_voice/model.py b/funasr/models/sense_voice/model.py
index 81feea9..70cd02e 100644
--- a/funasr/models/sense_voice/model.py
+++ b/funasr/models/sense_voice/model.py
@@ -19,7 +19,7 @@
 
 
 from funasr.models.paraformer.search import Hypothesis
-from funasr.models.sense_voice.utils.ctc_alignment import ctc_forced_align
+from .utils.ctc_alignment import ctc_forced_align
 
 
 class SinusoidalPositionEncoder(torch.nn.Module):
@@ -557,7 +557,7 @@
     ):
         """Embed positions in tensor."""
         maxlen = xs_pad.shape[1]
-        masks = sequence_mask(ilens, maxlen = maxlen, device=ilens.device)[:, None, :]
+        masks = sequence_mask(ilens, maxlen=maxlen, device=ilens.device)[:, None, :]
 
         xs_pad *= self.output_size() ** 0.5
 
@@ -916,27 +916,28 @@
 
             if output_timestamp:
                 from itertools import groupby
+
                 timestamp = []
                 tokens = tokenizer.text2tokens(text)[4:]
-                logits_speech = self.ctc.softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :]
+                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
+                logits_speech[pred == self.blank_id, self.blank_id] = 0
                 align = ctc_forced_align(
                     logits_speech.unsqueeze(0).float(),
                     torch.Tensor(token_int[4:]).unsqueeze(0).long().to(logits_speech.device),
-                    (encoder_out_lens-4).long(),
-                    torch.tensor(len(token_int)-4).unsqueeze(0).long().to(logits_speech.device),
+                    (encoder_out_lens - 4).long(),
+                    torch.tensor(len(token_int) - 4).unsqueeze(0).long().to(logits_speech.device),
                     ignore_id=self.ignore_id,
                 )
-                pred = groupby(align[0, :encoder_out_lens[0]])
+                pred = groupby(align[0, : encoder_out_lens[0]])
                 _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)
+                        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
@@ -952,21 +953,22 @@
         timestamp_new = []
         for i, t in enumerate(timestamp):
             word, start, end = t
-            if word == '鈻�':
+            if word == "鈻�":
                 continue
             if i == 0:
                 # timestamp_new.append([word, start, end])
-                timestamp_new.append([int(start*1000), int(end*1000)])
+                timestamp_new.append([int(start * 1000), int(end * 1000)])
             elif word.startswith("鈻�") or len(word) == 1 or not word[1].isalpha():
                 word = word[1:]
                 # timestamp_new.append([word, start, end])
-                timestamp_new.append([int(start*1000), int(end*1000)])
+                timestamp_new.append([int(start * 1000), int(end * 1000)])
             else:
                 # timestamp_new[-1][0] += word
-                timestamp_new[-1][1] = int(end*1000)
+                timestamp_new[-1][1] = int(end * 1000)
         return timestamp_new
+
     def export(self, **kwargs):
-        from export_meta import export_rebuild_model
+        from .export_meta import export_rebuild_model
 
         if "max_seq_len" not in kwargs:
             kwargs["max_seq_len"] = 512
@@ -974,4 +976,3 @@
         return models
 
         return results, meta_data
-

--
Gitblit v1.9.1