From 7759ab5febbccdeb929c34a5d5dd4b5387bfd83c Mon Sep 17 00:00:00 2001
From: Yuming Zhang <736099538@qq.com>
Date: 星期五, 23 二月 2024 18:28:26 +0800
Subject: [PATCH] fix bug: 模型初始化可传入参数disable_pbar=True (#1387)
---
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