纠错算法

  1. # @Author:sunshine
  2. # @Time  : 2020/5/12 上午9:17
  4. import json
  5. from bert4keras.tokenizers import load_vocab, Tokenizer
  6. from bert4keras.models import build_transformer_model
  7. from bert4keras.snippets import DataGenerator, sequence_padding
  8. from bert4keras.optimizers import AdaFactor
  9. from keras.layers import Lambda
  10. from keras.models import Model
  11. from keras.callbacks import Callback
  12. import keras.backend as K
  13. import numpy as np
  14. from tqdm import tqdm
  16. max_len = 64
  17. config_path = '/home/chenbing/pretrain_models/bert/chinese_L-12_H-768_A-12/bert_config.json'
  18. checkpoint_path = '/home/chenbing/pretrain_models/bert/chinese_L-12_H-768_A-12/bert_model.ckpt'
  19. vocab_path = '/home/chenbing/pretrain_models/bert/chinese_L-12_H-768_A-12/vocab.txt'
  21. train_data = json.load(open('data/train_data.json', 'r', encoding='utf-8'))
  22. valid_data = json.load(open('data/valid_data.json', 'r', encoding='utf-8'))
  24. # 加载精简词表
  25. token_dict, keep_words = load_vocab(
  26.     dict_path=vocab_path,
  27.     simplified=True,
  28.     startswith=['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]']
  29. )
  31. tokenizer = Tokenizer(token_dict, do_lower_case=True)
  34. class MyDataGenerator(DataGenerator):
  35.     def __iter__(self, random=True):
  36.         """
  37.         单条样本格式: [cls]错误词汇[sep][mask][mask]..[sep]
  38.         :param random:
  39.         :return:
  40.         """
  41.         batch_tokens_ids, batch_segment_ids, batch_right_token_ids = [], [], []
  42.         for is_end, D in self.sample(random):
  43.             wrong, right = D
  44.             right_token_ids, _ = tokenizer.encode(first_text=right)
  45.             wrong_token_ids, _ = tokenizer.encode(first_text=wrong)
  47.             token_ids = wrong_token_ids
  48.             token_ids += [tokenizer._token_mask_id] * max_len
  49.             token_ids += [tokenizer._token_end_id]
  51.             segemnt_ids = [0] * len(token_ids)
  53.             batch_tokens_ids.append(token_ids)
  54.             batch_segment_ids.append(segemnt_ids)
  55.             batch_right_token_ids.append(right_token_ids[1:])
  57.             if len(batch_tokens_ids) == self.batch_size or is_end:
  58.                 batch_tokens_ids = sequence_padding(batch_tokens_ids)
  59.                 batch_segment_ids = sequence_padding(batch_segment_ids)
  60.                 batch_right_token_ids = sequence_padding(batch_right_token_ids, max_len)
  62.                 yield [batch_tokens_ids, batch_segment_ids], batch_right_token_ids
  63.                 batch_tokens_ids, batch_segment_ids, batch_right_token_ids = [], [], []
  66. # 构建模型
  67. bert_model = build_transformer_model(
  68.     config_path=config_path,
  69.     checkpoint_path=checkpoint_path,
  70.     with_mlm=True,
  71.     keep_tokens=keep_words
  72. )
  74. output = Lambda(lambda x: x[:, 1:max_len + 1])(bert_model.output)
  75. model = Model(bert_model.input, output)
  78. def masked_cross_entropy(y_true, y_pred):
  79.     """交叉熵作为loss，并mask掉padding部分的预测
  80.     """
  81.     y_true = K.reshape(y_true, [K.shape(y_true)[0], -1])
  82.     y_mask = K.cast(K.not_equal(y_true, 0), K.floatx())
  83.     cross_entropy = K.sparse_categorical_crossentropy(y_true, y_pred)
  84.     cross_entropy = K.sum(cross_entropy * y_mask) / K.sum(y_mask)
  85.     return cross_entropy
  88. model.compile(loss=masked_cross_entropy, optimizer=AdaFactor(learning_rate=1e-3))
  89. model.summary()
  92. def ge_answer(wrong):
  93.     """
  94.     解码
  95.     :param wrong:
  96.     :return:
  97.     """
  98.     wrong_token_ids, _ = tokenizer.encode(wrong)
  99.     token_ids = wrong_token_ids + [tokenizer._token_mask_id] * max_len + [tokenizer._token_end_id]
  100.     segemnt_ids = [0] * len(token_ids)
  101.     probas = model.predict([np.array([token_ids]), np.array([segemnt_ids])])[0]
  102.     proba_ids = probas.argmax(axis=1)
  103.     useful_index = proba_ids[np.where(proba_ids != 3)]
  104.     if any(useful_index):
  105.         answer = tokenizer.decode(useful_index)
  106.     else:
  107.         answer = tokenizer.decode(proba_ids[:len(wrong)])
  108.     return answer
  111. def evalute(valid_data):
  112.     X, Y = 1e-10, 1e-10
  113.     for item in tqdm(valid_data):
  114.         wrong, right = item
  115.         pred = ge_answer(wrong)
  116.         X += pred == right
  117.         Y += 1
  118.     precision = X / Y
  119.     return precision
  122. class Evaluator(Callback):
  123.     def __init__(self):
  124.         self.lowest = 1e10
  125.         self.f1 = 1e-10
  127.     def on_epoch_end(self, epoch, logs=None):
  128.         if logs['loss'] <= self.lowest:
  129.             self.lowest = logs['loss']
  130.             model.save('models/best_mlm_model.h5')
  133. if __name__ == '__main__':
  134.     # 训练模型
  135.     # evaluator = Evaluator()
  136.     # train_generator = MyDataGenerator(train_data, batch_size=8)
  137.     #
  138.     # model.fit_generator(
  139.     #     train_generator.forfit(),
  140.     #     steps_per_epoch=len(train_generator),
  141.     #     epochs=10,
  142.     #     callbacks=[evaluator]
  143.     # )
  145.     # predict
  146.     model.load_weights('models/best_mlm_model.h5')
  147.     wrong = '追风少俊年王俊凯'
  148.     result = ge_answer(wrong)
  149.     print(result)
  1. # @Author:sunshine
  2. # @Time  : 2020/5/12 下午1:27
  4. import json
  5. from bert4keras.tokenizers import load_vocab, Tokenizer
  6. from bert4keras.models import build_transformer_model
  7. from bert4keras.snippets import DataGenerator, sequence_padding
  8. from bert4keras.optimizers import AdaFactor
  9. from keras.callbacks import Callback
  10. import keras.backend as K
  11. from tqdm import tqdm
  13. max_len = 64
  14. config_path = '/home/chenbing/pretrain_models/bert/chinese_L-12_H-768_A-12/bert_config.json'
  15. checkpoint_path = '/home/chenbing/pretrain_models/bert/chinese_L-12_H-768_A-12/bert_model.ckpt'
  16. vocab_path = '/home/chenbing/pretrain_models/bert/chinese_L-12_H-768_A-12/vocab.txt'
  18. train_data = json.load(open('data/train_data.json', 'r', encoding='utf-8'))
  19. valid_data = json.load(open('data/valid_data.json', 'r', encoding='utf-8'))
  21. # 加载精简词表
  22. token_dict, keep_words = load_vocab(
  23.     dict_path=vocab_path,
  24.     simplified=True,
  25.     startswith=['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]']
  26. )
  28. tokenizer = Tokenizer(token_dict, do_lower_case=True)
  31. class MyDataGenerator(DataGenerator):
  32.     def __iter__(self, random=True):
  33.         """
  34.         单条样本格式: [cls]错误词汇[sep][mask][mask]..[sep]
  35.         :param random:
  36.         :return:
  37.         """
  38.         batch_tokens_ids, batch_segment_ids = [], []
  39.         for is_end, D in self.sample(random):
  40.             wrong, right = D
  41.             # segment_ids也作为mask输入
  42.             token_ids, segment_ids = tokenizer.encode(first_text=wrong, second_text=right, max_length=max_len * 2)
  44.             batch_tokens_ids.append(token_ids)
  45.             batch_segment_ids.append(segment_ids)
  47.             if len(batch_tokens_ids) == self.batch_size or is_end:
  48.                 batch_tokens_ids = sequence_padding(batch_tokens_ids)
  49.                 batch_segment_ids = sequence_padding(batch_segment_ids)
  51.                 yield [batch_tokens_ids, batch_segment_ids], None
  52.                 batch_tokens_ids, batch_segment_ids = [], []
  55. # 构建模型
  56. model = build_transformer_model(
  57.     config_path=config_path,
  58.     checkpoint_path=checkpoint_path,
  59.     application='unilm',
  60.     keep_tokens=keep_words
  61. )
  63. y_true = model.input[0][:, 1:]
  64. y_mask = model.input[1][:, 1:]
  65. y_pred = model.output[:, :-1]
  67. cross_entropy = K.sparse_categorical_crossentropy(y_true, y_pred)
  68. cross_entropy = K.sum(cross_entropy * y_mask) / K.sum(y_mask)
  70. model.add_loss(cross_entropy)
  71. model.compile(optimizer=AdaFactor(learning_rate=1e-3))
  72. model.summary()
  75. def ge_answer(wrong):
  76.     """
  77.     解码
  78.     :param wrong:
  79.     :return:
  80.     """
  81.     wrong_token_ids, _ = tokenizer.encode(wrong)
  82.     token_ids = wrong_token_ids + [tokenizer._token_mask_id] * max_len + [tokenizer._token_end_id]
  83.     segemnt_ids = [0] * len(token_ids)
  84.     probas = model.predict([np.array([token_ids]), np.array([segemnt_ids])])[0]
  85.     proba_ids = probas.argmax(axis=1)
  86.     useful_index = proba_ids[np.where(proba_ids != 3)]
  87.     if any(useful_index):
  88.         answer = tokenizer.decode(useful_index)
  89.     else:
  90.         answer = tokenizer.decode(proba_ids[:len(wrong)])
  91.     return answer
  94. def evalute(valid_data):
  95.     X, Y = 1e-10, 1e-10
  96.     for item in tqdm(valid_data):
  97.         wrong, right = item
  98.         pred = ge_answer(wrong)
  99.         X += pred == right
  100.         Y += 1
  101.     precision = X / Y
  102.     return precision
  105. class Evaluator(Callback):
  106.     def __init__(self):
  107.         self.lowest = 1e10
  109.     def on_epoch_end(self, epoch, logs=None):
  110.         if logs['loss'] <= self.lowest:
  111.             self.lowest = logs['loss']
  112.             model.save('models/best_seq2seq_model.h5')
  115. if __name__ == '__main__':
  116.     evaluator = Evaluator()
  117.     train_generator = MyDataGenerator(train_data, batch_size=8)
  119.     model.fit_generator(
  120.         train_generator.forfit(),
  121.         steps_per_epoch=len(train_generator),
  122.         epochs=10,
  123.         callbacks=[evaluator]
  124.     )