| @@ -6,9 +6,7 @@ from ...utils.constant import Tasks | |||||
| from ..base import Model, Tensor | from ..base import Model, Tensor | ||||
| from ..builder import MODELS | from ..builder import MODELS | ||||
| __all__ = [ | |||||
| 'StructBertForMaskedLM', 'VecoForMaskedLM', 'AliceMindBaseForMaskedLM' | |||||
| ] | |||||
| __all__ = ['StructBertForMaskedLM', 'VecoForMaskedLM'] | |||||
| class AliceMindBaseForMaskedLM(Model): | class AliceMindBaseForMaskedLM(Model): | ||||
| @@ -40,9 +38,13 @@ class AliceMindBaseForMaskedLM(Model): | |||||
| @MODELS.register_module(Tasks.fill_mask, module_name=r'sbert') | @MODELS.register_module(Tasks.fill_mask, module_name=r'sbert') | ||||
| class StructBertForMaskedLM(AliceMindBaseForMaskedLM): | class StructBertForMaskedLM(AliceMindBaseForMaskedLM): | ||||
| # The StructBert for MaskedLM uses the same underlying model structure | |||||
| # as the base model class. | |||||
| pass | pass | ||||
| @MODELS.register_module(Tasks.fill_mask, module_name=r'veco') | @MODELS.register_module(Tasks.fill_mask, module_name=r'veco') | ||||
| class VecoForMaskedLM(AliceMindBaseForMaskedLM): | class VecoForMaskedLM(AliceMindBaseForMaskedLM): | ||||
| # The Veco for MaskedLM uses the same underlying model structure | |||||
| # as the base model class. | |||||
| pass | pass | ||||
| @@ -1,7 +1,8 @@ | |||||
| from typing import Dict, Optional | |||||
| from typing import Dict, Optional, Union | |||||
| from modelscope.models import Model | 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.preprocessors import FillMaskPreprocessor | ||||
| from modelscope.utils.constant import Tasks | from modelscope.utils.constant import Tasks | ||||
| from ..base import Pipeline, Tensor | from ..base import Pipeline, Tensor | ||||
| @@ -15,7 +16,7 @@ __all__ = ['FillMaskPipeline'] | |||||
| class FillMaskPipeline(Pipeline): | class FillMaskPipeline(Pipeline): | ||||
| def __init__(self, | def __init__(self, | ||||
| model: AliceMindBaseForMaskedLM, | |||||
| model: Union[AliceMindBaseForMaskedLM, str], | |||||
| preprocessor: Optional[FillMaskPreprocessor] = None, | preprocessor: Optional[FillMaskPreprocessor] = None, | ||||
| **kwargs): | **kwargs): | ||||
| """use `model` and `preprocessor` to create a nlp fill mask pipeline for prediction | """use `model` and `preprocessor` to create a nlp fill mask pipeline for prediction | ||||
| @@ -180,7 +180,7 @@ class TextGenerationPreprocessor(Preprocessor): | |||||
| return {k: torch.tensor(v) for k, v in rst.items()} | 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): | class FillMaskPreprocessor(Preprocessor): | ||||
| def __init__(self, model_dir: str, *args, **kwargs): | def __init__(self, model_dir: str, *args, **kwargs): | ||||