本项目开源在中文文本纠错项目:
pycorrector,可支持macbert4csc模型,通过如下命令调用:
1from pycorrector.macbert.macbert_corrector import MacBertCorrector
2
3m = MacBertCorrector("shibing624/macbert4csc-base-chinese")
4
5i = m.correct('今天新情很好')
6print(i)
1import operator
2import torch
3from transformers import BertTokenizer, BertForMaskedLM
4device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5
6tokenizer = BertTokenizer.from_pretrained("shibing624/macbert4csc-base-chinese")
7model = BertForMaskedLM.from_pretrained("shibing624/macbert4csc-base-chinese")
8model.to(device)
9
10texts = ["今天新情很好", "你找到你最喜欢的工作,我也很高心。"]
11with torch.no_grad():
12 outputs = model(**tokenizer(texts, padding=True, return_tensors='pt').to(device))
13
14def get_errors(corrected_text, origin_text):
15 sub_details = []
16 for i, ori_char in enumerate(origin_text):
17 if ori_char in [' ', '“', '”', '‘', '’', '琊', '\n', '…', '—', '擤']:
18 # add unk word
19 corrected_text = corrected_text[:i] + ori_char + corrected_text[i:]
20 continue
21 if i >= len(corrected_text):
22 continue
23 if ori_char != corrected_text[i]:
24 if ori_char.lower() == corrected_text[i]:
25 # pass english upper char
26 corrected_text = corrected_text[:i] + ori_char + corrected_text[i + 1:]
27 continue
28 sub_details.append((ori_char, corrected_text[i], i, i + 1))
29 sub_details = sorted(sub_details, key=operator.itemgetter(2))
30 return corrected_text, sub_details
31
32result = []
33for ids, text in zip(outputs.logits, texts):
34 _text = tokenizer.decode(torch.argmax(ids, dim=-1), skip_special_tokens=True).replace(' ', '')
35 corrected_text = _text[:len(text)]
36 corrected_text, details = get_errors(corrected_text, text)
37 print(text, ' => ', corrected_text, details)
38 result.append((corrected_text, details))
39print(result)
1今天新情很好 => 今天心情很好 [('新', '心', 2, 3)]
2你找到你最喜欢的工作,我也很高心。 => 你找到你最喜欢的工作,我也很高兴。 [('心', '兴', 15, 16)]
macbert4csc-base-chinese
├── config.json
├── added_tokens.json
├── pytorch_model.bin
├── special_tokens_map.json
├── tokenizer_config.json
└── vocab.txt
1[
2 {
3 "id": "B2-4029-3",
4 "original_text": "晚间会听到嗓音,白天的时候大家都不会太在意,但是在睡觉的时候这嗓音成为大家的恶梦。",
5 "wrong_ids": [
6 5,
7 31
8 ],
9 "correct_text": "晚间会听到噪音,白天的时候大家都不会太在意,但是在睡觉的时候这噪音成为大家的恶梦。"
10 },
11]
1macbert4csc
2 ├── config.json
3 ├── pytorch_model.bin
4 ├── special_tokens_map.json
5 ├── tokenizer_config.json
6 └── vocab.txt
Here is an example of our pre-training task.
Except for the new pre-training task, we also incorporate the following techniques.
1@software{pycorrector,
2 author = {Xu Ming},
3 title = {pycorrector: Text Error Correction Tool},
4 year = {2021},
5 url = {https://github.com/shibing624/pycorrector},
6}