OpenCompass/opencompass/models/ola/model/speech_encoder/builder.py
2025-03-04 23:10:00 +08:00

14 lines
531 B
Python

from .speech_encoder import WhisperWrappedEncoder, DualWrappedEncoder
import torch.nn as nn
def build_speech_encoder(config):
speech_encoder_type = getattr(config, 'speech_encoder_type', None)
if "whisper" in speech_encoder_type.lower():
return WhisperWrappedEncoder.load(config)
elif "dual" in speech_encoder_type.lower():
return DualWrappedEncoder(config)
elif "none" in speech_encoder_type.lower():
return None
raise ValueError(f'Unknown speech encoder: {speech_encoder_type}')