From b049f17f547320c652289bb5b46dedcf8f1e7f6e Mon Sep 17 00:00:00 2001
From: suluyan <suluyan.sly@alibaba-inc.com>
Date: Tue, 21 Jun 2022 15:52:22 +0800
Subject: [PATCH] fix comments

---
 modelscope/models/nlp/masked_language_model.py | 8 +++++---
 modelscope/pipelines/nlp/fill_mask_pipeline.py | 7 ++++---
 modelscope/preprocessors/nlp.py                | 2 +-
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/modelscope/models/nlp/masked_language_model.py b/modelscope/models/nlp/masked_language_model.py
index 848d7484..514c72c7 100644
--- a/modelscope/models/nlp/masked_language_model.py
+++ b/modelscope/models/nlp/masked_language_model.py
@@ -6,9 +6,7 @@ from ...utils.constant import Tasks
 from ..base import Model, Tensor
 from ..builder import MODELS
 
-__all__ = [
-    'StructBertForMaskedLM', 'VecoForMaskedLM', 'AliceMindBaseForMaskedLM'
-]
+__all__ = ['StructBertForMaskedLM', 'VecoForMaskedLM']
 
 
 class AliceMindBaseForMaskedLM(Model):
@@ -40,9 +38,13 @@ class AliceMindBaseForMaskedLM(Model):
 
 @MODELS.register_module(Tasks.fill_mask, module_name=r'sbert')
 class StructBertForMaskedLM(AliceMindBaseForMaskedLM):
+    # The StructBert for MaskedLM uses the same underlying model structure
+    # as the base model class.
     pass
 
 
 @MODELS.register_module(Tasks.fill_mask, module_name=r'veco')
 class VecoForMaskedLM(AliceMindBaseForMaskedLM):
+    # The Veco for MaskedLM uses the same underlying model structure
+    # as the base model class.
     pass
diff --git a/modelscope/pipelines/nlp/fill_mask_pipeline.py b/modelscope/pipelines/nlp/fill_mask_pipeline.py
index abe5b5b5..d7c1d456 100644
--- a/modelscope/pipelines/nlp/fill_mask_pipeline.py
+++ b/modelscope/pipelines/nlp/fill_mask_pipeline.py
@@ -1,7 +1,8 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 from modelscope.models import Model
-from modelscope.models.nlp import AliceMindBaseForMaskedLM
+from modelscope.models.nlp.masked_language_model import \
+    AliceMindBaseForMaskedLM
 from modelscope.preprocessors import FillMaskPreprocessor
 from modelscope.utils.constant import Tasks
 from ..base import Pipeline, Tensor
@@ -15,7 +16,7 @@ __all__ = ['FillMaskPipeline']
 class FillMaskPipeline(Pipeline):
 
     def __init__(self,
-                 model: AliceMindBaseForMaskedLM,
+                 model: Union[AliceMindBaseForMaskedLM, str],
                  preprocessor: Optional[FillMaskPreprocessor] = None,
                  **kwargs):
         """use `model` and `preprocessor` to create a nlp fill mask pipeline for prediction
diff --git a/modelscope/preprocessors/nlp.py b/modelscope/preprocessors/nlp.py
index c2f72292..fc910e98 100644
--- a/modelscope/preprocessors/nlp.py
+++ b/modelscope/preprocessors/nlp.py
@@ -180,7 +180,7 @@ class TextGenerationPreprocessor(Preprocessor):
         return {k: torch.tensor(v) for k, v in rst.items()}
 
 
-@PREPROCESSORS.register_module(Fields.nlp, module_name=r'sbert')
+@PREPROCESSORS.register_module(Fields.nlp)
 class FillMaskPreprocessor(Preprocessor):
 
     def __init__(self, model_dir: str, *args, **kwargs):