From bb2113434afe6b8d597256738abc7803b97c104c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=98=89=E6=B8=8A?= <wangjiaming.wjm@alibaba-inc.com>
Date: Mon, 24 Apr 2023 16:05:42 +0800
Subject: [PATCH] update

---
 funasr/models/target_delay_transformer.py |  5 +---
 funasr/models/vad_realtime_transformer.py |  3 +--
 funasr/train/abs_model.py                 | 28 +++--------------------
 3 files changed, 5 insertions(+), 31 deletions(-)

diff --git a/funasr/models/target_delay_transformer.py b/funasr/models/target_delay_transformer.py
index 84a2e6c99..683b804fc 100644
--- a/funasr/models/target_delay_transformer.py
+++ b/funasr/models/target_delay_transformer.py
@@ -6,13 +6,10 @@ import torch
 import torch.nn as nn
 
 from funasr.modules.embedding import SinusoidalPositionEncoder
-#from funasr.models.encoder.transformer_encoder import TransformerEncoder as Encoder
 from funasr.models.encoder.sanm_encoder import SANMEncoder as Encoder
-#from funasr.modules.mask import subsequent_n_mask
-from funasr.train.abs_model import AbsPunctuation
 
 
-class TargetDelayTransformer(AbsPunctuation):
+class TargetDelayTransformer(torch.nn.Module):
 
     def __init__(
         self,
diff --git a/funasr/models/vad_realtime_transformer.py b/funasr/models/vad_realtime_transformer.py
index 66f7fadeb..591ad4186 100644
--- a/funasr/models/vad_realtime_transformer.py
+++ b/funasr/models/vad_realtime_transformer.py
@@ -7,10 +7,9 @@ import torch.nn as nn
 
 from funasr.modules.embedding import SinusoidalPositionEncoder
 from funasr.models.encoder.sanm_encoder import SANMVadEncoder as Encoder
-from funasr.train.abs_model import AbsPunctuation
 
 
-class VadRealtimeTransformer(AbsPunctuation):
+class VadRealtimeTransformer(torch.nn.Module):
 
     def __init__(
         self,
diff --git a/funasr/train/abs_model.py b/funasr/train/abs_model.py
index 1c7ff3d33..026140b7d 100644
--- a/funasr/train/abs_model.py
+++ b/funasr/train/abs_model.py
@@ -12,34 +12,12 @@ from typeguard import check_argument_types
 
 from funasr.modules.nets_utils import make_pad_mask
 from funasr.torch_utils.device_funcs import force_gatherable
-from funasr.train.abs_espnet_model import AbsESPnetModel
-
-from funasr.modules.scorers.scorer_interface import BatchScorerInterface
+from funasr.models.base_model import FunASRModel
 
 
-class AbsPunctuation(torch.nn.Module, BatchScorerInterface, ABC):
-    """The abstract class
-
-    To share the loss calculation way among different models,
-    We uses delegate pattern here:
-    The instance of this class should be passed to "LanguageModel"
-
-    This "model" is one of mediator objects for "Task" class.
-
-    """
-
-    @abstractmethod
-    def forward(self, input: torch.Tensor, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        raise NotImplementedError
-
-    @abstractmethod
-    def with_vad(self) -> bool:
-        raise NotImplementedError
-
-
-class PunctuationModel(AbsESPnetModel):
+class PunctuationModel(FunASRModel):
     
-    def __init__(self, punc_model: AbsPunctuation, vocab_size: int, ignore_id: int = 0, punc_weight: list = None):
+    def __init__(self, punc_model: torch.nn.Module, vocab_size: int, ignore_id: int = 0, punc_weight: list = None):
         assert check_argument_types()
         super().__init__()
         self.punc_model = punc_model