Browse Source

fix comments

master
suluyan 3 years ago
parent
commit
b049f17f54
3 changed files with 10 additions and 7 deletions
  1. +5
    -3
      modelscope/models/nlp/masked_language_model.py
  2. +4
    -3
      modelscope/pipelines/nlp/fill_mask_pipeline.py
  3. +1
    -1
      modelscope/preprocessors/nlp.py

+ 5
- 3
modelscope/models/nlp/masked_language_model.py View File

@@ -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

+ 4
- 3
modelscope/pipelines/nlp/fill_mask_pipeline.py View File

@@ -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


+ 1
- 1
modelscope/preprocessors/nlp.py View File

@@ -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):


Loading…
Cancel
Save