| 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 |
|
|
