[ssl] init ssl model (#2595)

* init ssl model

* fix whisper ut

* try to fix ut

wenet/ssl/init_model.py

@@ -0,0 +1,19 @@
+from wenet.ssl.bestrq.bestrq_model import BestRQModel
+from wenet.ssl.wav2vec2.wav2vec2_model import Wav2vec2Model
+from wenet.ssl.w2vbert.w2vbert_model import W2VBERTModel
+
+WENET_SSL_MODEL_CLASS = {
+    "w2vbert_model": W2VBERTModel,
+    "wav2vec_model": Wav2vec2Model,
+    "bestrq_model": BestRQModel
+}
+
+
+def init_model(configs, encoder):
+
+    assert 'model' in configs
+    model_type = configs['model']
+    assert model_type in WENET_SSL_MODEL_CLASS.keys()
+    model = WENET_SSL_MODEL_CLASS[model_type](encoder=encoder,
+                                              **configs['model_conf'])
+    return model

wenet/utils/init_model.py

@@ -23,6 +23,7 @@
 from wenet.paraformer.paraformer import Paraformer, Predictor
 from wenet.LLM.causallm_model import CausalLM
 from wenet.LLM.decoder import DecoderOnly
+from wenet.ssl.init_model import WENET_SSL_MODEL_CLASS
 from wenet.transducer.joint import TransducerJoint
 from wenet.transducer.predictor import (ConvPredictor, EmbeddingPredictor,
                                         RNNPredictor)
@@ -42,7 +43,6 @@
 from wenet.utils.cmvn import load_cmvn
 from wenet.utils.checkpoint import load_checkpoint, load_trained_modules
 
-
 WENET_ENCODER_CLASSES = {
     "transformer": TransformerEncoder,
     "conformer": ConformerEncoder,
@@ -156,6 +156,9 @@ def init_speech_model(args, configs):
             special_tokens=configs.get('tokenizer_conf',
                                        {}).get('special_tokens', None),
         )
+    elif model_type in WENET_SSL_MODEL_CLASS.keys():
+        from wenet.ssl.init_model import init_model as init_ssl_model
+        model = init_ssl_model(configs, encoder)
     else:
         model = WENET_MODEL_CLASSES[model_type](
             vocab_size=vocab_size,
@@ -187,6 +190,7 @@ def init_causal_llm(configs):
 def init_model(args, configs):
 
     model_type = configs.get('model', 'asr_model')
+    configs['model'] = model_type
     if model_type == 'causal_lm':
         model, configs = init_causal_llm(configs)
     else: