forked from shibing624/pycorrector
-
Notifications
You must be signed in to change notification settings - Fork 0
/
corrector.py
319 lines (294 loc) · 11.9 KB
/
corrector.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
# -*- coding: utf-8 -*-
"""
@author:XuMing(xuming624@qq.com)
@description: corrector with pinyin and stroke
"""
import operator
import os
from codecs import open
from loguru import logger
import pypinyin
from pycorrector import config
from pycorrector.detector import Detector, ErrorType
from pycorrector.utils.math_utils import edit_distance_word
from pycorrector.utils.text_utils import is_chinese_string
from pycorrector.utils.tokenizer import segment, split_2_short_text
class Corrector(Detector):
def __init__(
self,
common_char_path=config.common_char_path,
same_pinyin_path=config.same_pinyin_path,
same_stroke_path=config.same_stroke_path,
language_model_path=config.language_model_path,
word_freq_path=config.word_freq_path,
custom_word_freq_path='',
custom_confusion_path='',
person_name_path=config.person_name_path,
place_name_path=config.place_name_path,
stopwords_path=config.stopwords_path,
proper_name_path=config.proper_name_path,
stroke_path=config.stroke_path
):
super(Corrector, self).__init__(
language_model_path=language_model_path,
word_freq_path=word_freq_path,
custom_word_freq_path=custom_word_freq_path,
custom_confusion_path=custom_confusion_path,
person_name_path=person_name_path,
place_name_path=place_name_path,
stopwords_path=stopwords_path,
proper_name_path=proper_name_path,
stroke_path=stroke_path
)
self.name = 'corrector'
self.common_char_path = common_char_path
self.same_pinyin_text_path = same_pinyin_path
self.same_stroke_text_path = same_stroke_path
self.initialized_corrector = False
self.cn_char_set = None
self.same_pinyin = None
self.same_stroke = None
@staticmethod
def load_set_file(path):
words = set()
with open(path, 'r', encoding='utf-8') as f:
for w in f:
w = w.strip()
if w.startswith('#'):
continue
if w:
words.add(w)
return words
@staticmethod
def load_same_pinyin(path, sep='\t'):
"""
加载同音字
:param path:
:param sep:
:return:
"""
result = dict()
if not os.path.exists(path):
logger.warn("file not exists:" + path)
return result
with open(path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if line.startswith('#'):
continue
parts = line.split(sep)
if parts and len(parts) > 2:
key_char = parts[0]
same_pron_same_tone = set(list(parts[1]))
same_pron_diff_tone = set(list(parts[2]))
value = same_pron_same_tone.union(same_pron_diff_tone)
if key_char and value:
result[key_char] = value
return result
@staticmethod
def load_same_stroke(path, sep='\t'):
"""
加载形似字
:param path:
:param sep:
:return:
"""
result = dict()
if not os.path.exists(path):
logger.warn("file not exists:" + path)
return result
with open(path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
if line.startswith('#'):
continue
parts = line.split(sep)
if parts and len(parts) > 1:
for i, c in enumerate(parts):
exist = result.get(c, set())
current = set(list(parts[:i] + parts[i + 1:]))
result[c] = exist.union(current)
return result
def _initialize_corrector(self):
# chinese common char
self.cn_char_set = self.load_set_file(self.common_char_path)
# same pinyin
self.same_pinyin = self.load_same_pinyin(self.same_pinyin_text_path)
# same stroke
self.same_stroke = self.load_same_stroke(self.same_stroke_text_path)
self.initialized_corrector = True
def check_corrector_initialized(self):
if not self.initialized_corrector:
self._initialize_corrector()
def get_same_pinyin(self, char):
"""
取同音字
:param char:
:return:
"""
self.check_corrector_initialized()
return self.same_pinyin.get(char, set())
def get_same_stroke(self, char):
"""
取形似字
:param char:
:return:
"""
self.check_corrector_initialized()
return self.same_stroke.get(char, set())
def known(self, words):
"""
取得词序列中属于常用词部分
:param words:
:return:
"""
self.check_detector_initialized()
return set(word for word in words if word in self.word_freq)
def _confusion_char_set(self, c):
return self.get_same_pinyin(c).union(self.get_same_stroke(c))
def _confusion_word_set(self, word):
confusion_word_set = set()
candidate_words = list(self.known(edit_distance_word(word, self.cn_char_set)))
for candidate_word in candidate_words:
if pypinyin.lazy_pinyin(candidate_word) == pypinyin.lazy_pinyin(word):
# same pinyin
confusion_word_set.add(candidate_word)
return confusion_word_set
def _confusion_custom_set(self, word):
confusion_word_set = set()
if word in self.custom_confusion:
confusion_word_set = {self.custom_confusion[word]}
return confusion_word_set
def generate_items(self, word, fragment=1):
"""
生成纠错候选集
:param word:
:param fragment: 分段
:return:
"""
self.check_corrector_initialized()
# 1字
candidates_1 = []
# 2字
candidates_2 = []
# 多于2字
candidates_3 = []
# same pinyin word
candidates_1.extend(self._confusion_word_set(word))
# custom confusion word
candidates_1.extend(self._confusion_custom_set(word))
# get similarity char
if len(word) == 1:
# sim one char
confusion = [i for i in self._confusion_char_set(word[0]) if i]
candidates_1.extend(confusion)
if len(word) == 2:
# sim first char
confusion_first = [i for i in self._confusion_char_set(word[0]) if i]
candidates_2.extend([i + word[1] for i in confusion_first])
# sim last char
confusion_last = [i for i in self._confusion_char_set(word[1]) if i]
candidates_2.extend([word[0] + i for i in confusion_last])
# both change, sim char
candidates_2.extend([i + j for i in confusion_first for j in confusion_last if i + j])
# sim word
# candidates_2.extend([i for i in self._confusion_word_set(word) if i])
if len(word) > 2:
# sim mid char
confusion = [word[0] + i + word[2:] for i in self._confusion_char_set(word[1])]
candidates_3.extend(confusion)
# sim first word
confusion_word = [i + word[-1] for i in self._confusion_word_set(word[:-1])]
candidates_3.extend(confusion_word)
# sim last word
confusion_word = [word[0] + i for i in self._confusion_word_set(word[1:])]
candidates_3.extend(confusion_word)
# add all confusion word list
confusion_word_set = set(candidates_1 + candidates_2 + candidates_3)
confusion_word_list = [item for item in confusion_word_set if is_chinese_string(item)]
confusion_sorted = sorted(confusion_word_list, key=lambda k: self.word_frequency(k), reverse=True)
return confusion_sorted[:len(confusion_word_list) // fragment + 1]
def get_lm_correct_item(self, cur_item, candidates, before_sent, after_sent, threshold=57, cut_type='char'):
"""
通过语言模型纠正字词错误
:param cur_item: 当前词
:param candidates: 候选词
:param before_sent: 前半部分句子
:param after_sent: 后半部分句子
:param threshold: ppl阈值, 原始字词替换后大于该ppl值则认为是错误
:param cut_type: 切词方式, 字粒度
:return: str, correct item, 正确的字词
"""
result = cur_item
if cur_item not in candidates:
candidates.append(cur_item)
ppl_scores = {i: self.ppl_score(segment(before_sent + i + after_sent, cut_type=cut_type)) for i in candidates}
sorted_ppl_scores = sorted(ppl_scores.items(), key=lambda d: d[1])
# 增加正确字词的修正范围,减少误纠
top_items = []
top_score = 0.0
for i, v in enumerate(sorted_ppl_scores):
v_word = v[0]
v_score = v[1]
if i == 0:
top_score = v_score
top_items.append(v_word)
# 通过阈值修正范围
elif v_score < top_score + threshold:
top_items.append(v_word)
else:
break
if cur_item not in top_items:
result = top_items[0]
return result
def correct(self, text, include_symbol=True, num_fragment=1, threshold=57, **kwargs):
"""
文本改错
改错逻辑:
1. 自定义混淆集
2. 专名错误
3. 字词错误
:param text: str, query 文本
:param include_symbol: bool, 是否包含标点符号
:param num_fragment: 纠错候选集分段数, 1 / (num_fragment + 1)
:param threshold: 语言模型纠错ppl阈值
:param kwargs: ...
:return: text (str)改正后的句子, list(wrong, right, begin_idx, end_idx)
"""
text_new = ''
details = []
self.check_corrector_initialized()
# 文本切分为句子
sentences = split_2_short_text(text, include_symbol=include_symbol)
for sentence, idx in sentences:
maybe_errors, proper_details = self.detect_sentence(sentence, idx, **kwargs)
for cur_item, begin_idx, end_idx, err_type in maybe_errors:
# 纠错,逐个处理
before_sent = sentence[:(begin_idx - idx)]
after_sent = sentence[(end_idx - idx):]
# 困惑集中指定的词,直接取结果
if err_type == ErrorType.confusion:
corrected_item = self.custom_confusion[cur_item]
elif err_type == ErrorType.proper:
# 专名错误 proper_details format: (error_word, corrected_word, begin_idx, end_idx)
corrected_item = [i[1] for i in proper_details if cur_item == i[0] and begin_idx == i[2]][0]
else:
# 字词错误,找所有可能正确的词
candidates = self.generate_items(cur_item, fragment=num_fragment)
if not candidates:
continue
corrected_item = self.get_lm_correct_item(
cur_item,
candidates,
before_sent,
after_sent,
threshold=threshold
)
# output
if corrected_item != cur_item:
sentence = before_sent + corrected_item + after_sent
detail_word = (cur_item, corrected_item, begin_idx, end_idx)
details.append(detail_word)
text_new += sentence
details = sorted(details, key=operator.itemgetter(2))
return text_new, details