funasr1.0 uniasr

modelscope · Jan 24, 2024 · 6cf6524 · 6cf6524
1 parent 3ead3cd
commit 6cf6524
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/funasr/models/uniasr/model.py b/funasr/models/uniasr/model.py
@@ -802,7 +802,12 @@ def init_beam_search(self,
         from funasr.models.uniasr.beam_search import BeamSearchScama
         from funasr.models.transformer.scorers.ctc import CTCPrefixScorer
         from funasr.models.transformer.scorers.length_bonus import LengthBonus
-
+
+        decoding_mode = kwargs.get("decoding_mode", "model1")
+        if decoding_mode == "model1":
+            decoder = self.decoder
+        else:
+            decoder = self.decoder2
         # 1. Build ASR model
         scorers = {}
 
@@ -813,7 +818,7 @@ def init_beam_search(self,
             )
         token_list = kwargs.get("token_list")
         scorers.update(
-            decoder=self.decoder,
+            decoder=decoder,
             length_bonus=LengthBonus(len(token_list)),
         )
 
@@ -830,7 +835,7 @@ def init_beam_search(self,
             length_bonus=kwargs.get("penalty", 0.0),
         )
         beam_search = BeamSearchScama(
-            beam_size=kwargs.get("beam_size", 2),
+            beam_size=kwargs.get("beam_size", 5),
             weights=weights,
             scorers=scorers,
             sos=self.sos,
@@ -866,7 +871,7 @@ def inference(self,
 
         if self.beam_search is None:
             logging.info("enable beam_search")
-            self.init_beam_search(**kwargs)
+            self.init_beam_search(decoding_mode=decoding_mode, **kwargs)
             self.nbest = kwargs.get("nbest", 1)
 
         meta_data = {}