diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 63af2ec4..78cc9219 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -45,6 +45,7 @@ class Pipelines(object): word_segmentation = 'word-segmentation' text_generation = 'text-generation' sentiment_analysis = 'sentiment-analysis' + fill_mask = 'fill-mask' # audio tasks sambert_hifigan_16k_tts = 'sambert-hifigan-16k-tts' diff --git a/modelscope/pipelines/nlp/fill_mask_pipeline.py b/modelscope/pipelines/nlp/fill_mask_pipeline.py index 291b1cdf..863d9a6d 100644 --- a/modelscope/pipelines/nlp/fill_mask_pipeline.py +++ b/modelscope/pipelines/nlp/fill_mask_pipeline.py @@ -1,5 +1,6 @@ from typing import Dict, Optional, Union +from modelscope.metainfo import Pipelines from modelscope.models import Model from modelscope.models.nlp.masked_language_model import \ AliceMindBaseForMaskedLM @@ -11,7 +12,7 @@ from ..builder import PIPELINES __all__ = ['FillMaskPipeline'] -@PIPELINES.register_module(Tasks.fill_mask, module_name=r'fill_mask') +@PIPELINES.register_module(Tasks.fill_mask, module_name=Pipelines.fill_mask) class FillMaskPipeline(Pipeline): def __init__(self,