From b52fc0ec2886c66c8b480a9f761a0af561afff5b Mon Sep 17 00:00:00 2001
From: Shi Xian <40013335+R1ckShi@users.noreply.github.com>
Date: 星期三, 28 二月 2024 14:39:54 +0800
Subject: [PATCH] Merge pull request #1401 from alibaba-damo-academy/dev_gzf
---
funasr/models/whisper/model.py | 8 ++++++--
1 files changed, 6 insertions(+), 2 deletions(-)
diff --git a/funasr/models/whisper/model.py b/funasr/models/whisper/model.py
index 6a6d47c..381a501 100644
--- a/funasr/models/whisper/model.py
+++ b/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__()
--
Gitblit v1.9.1