zhifu gao
2024-02-01 2ddfc27d5b69e5c1a38021554e97dae958328c20
funasr/models/whisper/model.py
@@ -10,6 +10,8 @@
from funasr.models.whisper.utils.decoding import detect_language as detect_language_function, decode as decode_function
from funasr.register import tables
@dataclass
class ModelDimensions:
@@ -128,6 +130,8 @@
        return x
@tables.register("encoder_classes", "WhisperEncoder")
class AudioEncoder(nn.Module):
    def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
        super().__init__()
@@ -158,7 +162,7 @@
        x = self.ln_post(x)
        return x
@tables.register("decoder_classes", "WhisperDecoder")
class TextDecoder(nn.Module):
    def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
        super().__init__()
@@ -193,7 +197,7 @@
        return logits
@tables.register("model_classes", "Whisper")
class Whisper(nn.Module):
    def __init__(self, dims: dict):
        super().__init__()