View on GitHub

pycorrector

pycorrector is a toolkit for text error correction. 文本纠错,Kenlm,Seq2Seq_Attention,BERT,MacBERT,ELECTRA,ERNIE,Transformer等模型实现,开箱即用。

MacBertMaskedLM For Correction

使用说明

快速加载

本项目是MacBERT的中文文本纠错模型,可支持BERT模型,可通过如下命令调用:

example: correct_demo.py

from pycorrector.macbert.macbert_corrector import MacBertCorrector

nlp = MacBertCorrector("shibing624/macbert4csc-base-chinese").macbert_correct

i = nlp('今天新情很好')
print(i)

当然,你也可使用官方的transformers库进行调用。

1.先pip安装transformers库:

pip install transformers>=4.1.1

2.使用以下示例执行:

import operator
import torch
from transformers import BertTokenizer, BertForMaskedLM
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = BertTokenizer.from_pretrained("shibing624/macbert4csc-base-chinese")
model = BertForMaskedLM.from_pretrained("shibing624/macbert4csc-base-chinese")
model = model.to(device)

texts = ["今天新情很好", "你找到你最喜欢的工作,我也很高心。"]
with torch.no_grad():
    outputs = model(**tokenizer(texts, padding=True, return_tensors='pt').to(device))

def get_errors(corrected_text, origin_text):
    sub_details = []
    for i, ori_char in enumerate(origin_text):
        if ori_char in [' ', '“', '”', '‘', '’', '琊', '\n', '…', '—', '擤']:
            # add unk word
            corrected_text = corrected_text[:i] + ori_char + corrected_text[i:]
            continue
        if i >= len(corrected_text):
            continue
        if ori_char != corrected_text[i]:
            if ori_char.lower() == corrected_text[i]:
                # pass english lower char
                corrected_text = corrected_text[:i] + ori_char + corrected_text[i + 1:]
                continue
            sub_details.append((ori_char, corrected_text[i], i, i + 1))
    sub_details = sorted(sub_details, key=operator.itemgetter(2))
    return corrected_text, sub_details

result = []
for ids, text in zip(outputs.logits, texts):
    _text = tokenizer.decode(torch.argmax(ids, dim=-1), skip_special_tokens=True).replace(' ', '')
    corrected_text = _text[:len(text)]
    corrected_text, details = get_errors(corrected_text, text)
    print(text, ' => ', corrected_text, details)
    result.append((corrected_text, details))
print(result)

output:

今天新情很好  =>  今天心情很好 [('新', '心', 2, 3)]
你找到你最喜欢的工作,我也很高心。  =>  你找到你最喜欢的工作,我也很高兴。 [('心', '兴', 15, 16)]

模型文件组成:

macbert4csc-base-chinese
    ├── config.json
    ├── added_tokens.json
    ├── pytorch_model.bin
    ├── special_tokens_map.json
    ├── tokenizer_config.json
    └── vocab.txt

Evaluate

提供评估脚本pycorrector/utils/eval.py,该脚本有两个功能:

执行该评估脚本后,

shibing624/macbert4csc-base-chinese 模型在corpus500纠错效果评估如下:

规则方法(加入自定义混淆集)在corpus500纠错效果评估如下:

shibing624/macbert4csc-base-chinese 在 SIGHAN2015 测试集纠错效果评估如下:

由于训练使用的数据使用了SIGHAN2015的训练集(复现paper),在SIGHAN2015的测试集上达到SOTA水平。

评估case

纠错结果除部分英文大小写问题外,在sighan15上达到了SOTA水平。

训练

安装依赖

pip install transformers>=4.1.1 pytorch-lightning>=1.1.2 torch>=1.7.0 

训练数据集

toy数据集(约1千条)

cd macbert
python preprocess.py

得到toy数据集文件:

macbert/output
|-- dev.json
|-- test.json
`-- train.json

SIGHAN+Wang271K中文纠错数据集

数据集 语料 下载链接 压缩包大小
SIGHAN+Wang271K中文纠错数据集 SIGHAN+Wang271K(27万条) 百度网盘(密码01b9) 106M
原始SIGHAN数据集 SIGHAN13 14 15 官方csc.html 339K
原始Wang271K数据集 Wang271K Automatic-Corpus-Generation dimmywang提供 93M

SIGHAN+Wang271K中文纠错数据集,数据格式:

[
    {
        "id": "B2-4029-3",
        "original_text": "晚间会听到嗓音,白天的时候大家都不会太在意,但是在睡觉的时候这嗓音成为大家的恶梦。",
        "wrong_ids": [
            5,
            31
        ],
        "correct_text": "晚间会听到噪音,白天的时候大家都不会太在意,但是在睡觉的时候这噪音成为大家的恶梦。"
    }
]

下载SIGHAN+Wang271K中文纠错数据集,下载后新建output文件夹并放里面,文件位置同上。

自有数据集

把自己数据集标注好,保存为跟训练样本集一样的json格式,然后加载模型继续训练即可。

  1. 已有大量业务相关错误样本,主要标注错误位置(wrong_ids)和纠错后的句子(correct_text)
  2. 没有现成的错误样本,可以手动写脚本生成错误样本(original_text),根据音似、形似等特征把正确句子的指定位置(wrong_ids)字符改为错字

训练

python train.py

预测

output
└── macbert4csc
    ├── config.json
    ├── pytorch_model.bin
    ├── special_tokens_map.json
    ├── tokenizer_config.json
    └── vocab.txt

demo示例macbert_corrector.py:

python3 macbert_corrector.py

如果需要训练SoftMaskedBertModel,请参考https://github.com/gitabtion/BertBasedCorrectionModels