# coding=utf-8
# author: xusong <xusong28@jd.com>
# time: 2022/8/25 10:19

"""
未完成

https://github.com/shibing624/pycorrector/blob/6f1b6ea35c60700f08463950e11ae0963ed512b2/pycorrector/bert/bert_corrector.py#L30
"""
import register
from transformers import pipeline

bert_model_dir = "models/pretrain/"
model = pipeline('fill-mask',
                      model=bert_model_dir,
                      tokenizer=bert_model_dir)

predicts = model("这款连[MASK]裙真漂亮")
print(predicts)