游雁 il y a 2 ans
Parent
commit
952c63f696

+ 5 - 3
funasr/models/bicif_paraformer/model.py

@@ -300,9 +300,11 @@ class BiCifParaformer(Paraformer):
                 nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
             for nbest_idx, hyp in enumerate(nbest_hyps):
                 ibest_writer = None
-                if ibest_writer is None and kwargs.get("output_dir") is not None:
-                    writer = DatadirWriter(kwargs.get("output_dir"))
-                    ibest_writer = writer[f"{nbest_idx + 1}best_recog"]
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx+1}best_recog"]
+                    
                 # remove sos/eos and get results
                 last_pos = -1
                 if isinstance(hyp.yseq, list):

+ 5 - 3
funasr/models/contextual_paraformer/model.py

@@ -380,9 +380,11 @@ class ContextualParaformer(Paraformer):
                 nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
             for nbest_idx, hyp in enumerate(nbest_hyps):
                 ibest_writer = None
-                if ibest_writer is None and kwargs.get("output_dir") is not None:
-                    writer = DatadirWriter(kwargs.get("output_dir"))
-                    ibest_writer = writer[f"{nbest_idx + 1}best_recog"]
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
+    
                 # remove sos/eos and get results
                 last_pos = -1
                 if isinstance(hyp.yseq, list):

+ 4 - 3
funasr/models/fsmn_vad_streaming/model.py

@@ -625,9 +625,10 @@ class FsmnVADStreaming(nn.Module):
 			self.init_cache(cache)
 		
 		ibest_writer = None
-		if ibest_writer is None and kwargs.get("output_dir") is not None:
-			writer = DatadirWriter(kwargs.get("output_dir"))
-			ibest_writer = writer[f"{1}best_recog"]
+		if kwargs.get("output_dir") is not None:
+			if not hasattr(self, "writer"):
+				self.writer = DatadirWriter(kwargs.get("output_dir"))
+			ibest_writer = self.writer[f"{1}best_recog"]
 		
 		results = []
 		result_i = {"key": key[0], "value": segments}

+ 5 - 3
funasr/models/monotonic_aligner/model.py

@@ -183,9 +183,11 @@ class MonotonicAligner(torch.nn.Module):
         
         results = []
         ibest_writer = None
-        if ibest_writer is None and kwargs.get("output_dir") is not None:
-            writer = DatadirWriter(kwargs.get("output_dir"))
-            ibest_writer = writer["tp_res"]
+        if kwargs.get("output_dir") is not None:
+            if not hasattr(self, "writer"):
+                self.writer = DatadirWriter(kwargs.get("output_dir"))
+            ibest_writer = self.writer["tp_res"]
+
         for i, (us_alpha, us_peak, token_int) in enumerate(zip(us_alphas, us_peaks, text_token_int_list)):
             token = tokenizer.ids2tokens(token_int)
             timestamp_str, timestamp = ts_prediction_lfr6_standard(us_alpha[:encoder_out_lens[i] * 3],

+ 4 - 3
funasr/models/paraformer/model.py

@@ -514,9 +514,10 @@ class Paraformer(torch.nn.Module):
                 nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
             for nbest_idx, hyp in enumerate(nbest_hyps):
                 ibest_writer = None
-                if ibest_writer is None and kwargs.get("output_dir") is not None:
-                    writer = DatadirWriter(kwargs.get("output_dir"))
-                    ibest_writer = writer[f"{nbest_idx+1}best_recog"]
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx+1}best_recog"]
                 # remove sos/eos and get results
                 last_pos = -1
                 if isinstance(hyp.yseq, list):

+ 4 - 3
funasr/models/paraformer_streaming/model.py

@@ -556,11 +556,12 @@ class ParaformerStreaming(Paraformer):
             self.init_cache(cache, **kwargs)
         
         if kwargs.get("output_dir"):
-            writer = DatadirWriter(kwargs.get("output_dir"))
-            ibest_writer = writer[f"{1}best_recog"]
+            if not hasattr(self, "writer"):
+                self.writer = DatadirWriter(kwargs.get("output_dir"))
+            ibest_writer = self.writer[f"{1}best_recog"]
             ibest_writer["token"][key[0]] = " ".join(tokens)
             ibest_writer["text"][key[0]] = text_postprocessed
-        
+
         return result, meta_data
 
 

+ 5 - 3
funasr/models/seaco_paraformer/model.py

@@ -385,9 +385,11 @@ class SeacoParaformer(BiCifParaformer, Paraformer):
                 nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
             for nbest_idx, hyp in enumerate(nbest_hyps):
                 ibest_writer = None
-                if ibest_writer is None and kwargs.get("output_dir") is not None:
-                    writer = DatadirWriter(kwargs.get("output_dir"))
-                    ibest_writer = writer[f"{nbest_idx + 1}best_recog"]
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
+                    
                 # remove sos/eos and get results
                 last_pos = -1
                 if isinstance(hyp.yseq, list):

+ 4 - 3
funasr/models/transducer/model.py

@@ -488,9 +488,10 @@ class Transducer(torch.nn.Module):
 
             for nbest_idx, hyp in enumerate(nbest_hyps):
                 ibest_writer = None
-                if ibest_writer is None and kwargs.get("output_dir") is not None:
-                    writer = DatadirWriter(kwargs.get("output_dir"))
-                    ibest_writer = writer[f"{nbest_idx+1}best_recog"]
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
                 # remove sos/eos and get results
                 last_pos = -1
                 if isinstance(hyp.yseq, list):

+ 5 - 3
funasr/models/transformer/model.py

@@ -423,9 +423,11 @@ class Transformer(nn.Module):
 
             for nbest_idx, hyp in enumerate(nbest_hyps):
                 ibest_writer = None
-                if ibest_writer is None and kwargs.get("output_dir") is not None:
-                    writer = DatadirWriter(kwargs.get("output_dir"))
-                    ibest_writer = writer[f"{nbest_idx+1}best_recog"]
+                if kwargs.get("output_dir") is not None:
+                    if not hasattr(self, "writer"):
+                        self.writer = DatadirWriter(kwargs.get("output_dir"))
+                    ibest_writer = self.writer[f"{nbest_idx + 1}best_recog"]
+                    
                 # remove sos/eos and get results
                 last_pos = -1
                 if isinstance(hyp.yseq, list):