import torch import os from enum import IntEnum from random import choice import random import collections import time import logging import jieba jieba.setLogLevel(logging.CRITICAL) import re import numpy as np import mask PAD = 0 MaskedLMInstance = collections.namedtuple("MaskedLMInstance", ["index", "label"]) def map_to_numpy(data): return np.asarray(data) class PreTrainingDataset(): def __init__(self, tokenizer, max_seq_length, backend='python', max_predictions_per_seq: int = 80, do_whole_word_mask: bool = True): self.tokenizer = tokenizer self.max_seq_length = max_seq_length self.masked_lm_prob = 0.15 self.backend = backend self.do_whole_word_mask = do_whole_word_mask self.max_predictions_per_seq = max_predictions_per_seq self.vocab_words = list(tokenizer.vocab.keys()) self.rec = re.compile('[\u4E00-\u9FA5]') self.whole_rec = re.compile('##[\u4E00-\u9FA5]') self.mlm_p = 0.15 self.mlm_mask_p = 0.8 self.mlm_tamper_p = 0.05 self.mlm_maintain_p = 0.1 def tokenize(self, doc): temp = [] for d in doc: temp.append(self.tokenizer.tokenize(d)) return temp def create_training_instance(self, instance): is_next = 1 raw_text_list = self.get_new_segment(instance) tokens_a = raw_text_list assert len(tokens_a) == len(instance) # tokens_a, tokens_b, is_next = instance.get_values() # print(f'is_next label:{is_next}') # Create mapper tokens = [] original_tokens = [] segment_ids = [] tokens.append("[CLS]") original_tokens.append('[CLS]') segment_ids.append(0) for index, token in enumerate(tokens_a): tokens.append(token) original_tokens.append(instance[index]) segment_ids.append(0) tokens.append("[SEP]") original_tokens.append('[SEP]') segment_ids.append(0) # for token in tokens_b: # tokens.append(token) # segment_ids.append(1) # tokens.append("[SEP]") # segment_ids.append(1) # Get Masked LM predictions if self.backend == 'c++': output_tokens, masked_lm_output = mask.create_whole_masked_lm_predictions(tokens, original_tokens, self.vocab_words, self.tokenizer.vocab, self.max_predictions_per_seq, self.masked_lm_prob) elif self.backend == 'python': output_tokens, masked_lm_output = self.create_whole_masked_lm_predictions(tokens) # Convert to Ids input_ids = self.tokenizer.convert_tokens_to_ids(output_tokens) input_mask = [1] * len(input_ids) while len(input_ids) < self.max_seq_length: input_ids.append(PAD) segment_ids.append(PAD) input_mask.append(PAD) masked_lm_output.append(-1) return ([ map_to_numpy(input_ids), map_to_numpy(input_mask), map_to_numpy(segment_ids), map_to_numpy(masked_lm_output), map_to_numpy([is_next]) ]) def create_masked_lm_predictions(self, tokens): cand_indexes = [] for i, token in enumerate(tokens): if token == "[CLS]" or token == "[SEP]": continue if (self.do_whole_word_mask and len(cand_indexes) >= 1 and token.startswith("##")): cand_indexes[-1].append(i) else: cand_indexes.append([i]) # cand_indexes.append(i) random.shuffle(cand_indexes) output_tokens = list(tokens) num_to_predict = min( self.max_predictions_per_seq, max(1, int(round(len(tokens) * self.masked_lm_prob)))) masked_lms = [] covered_indexes = set() for index in cand_indexes: if len(masked_lms) >= num_to_predict: break if index in covered_indexes: continue covered_indexes.add(index) masked_token = None # 80% mask if random.random() < 0.8: masked_token = "[MASK]" else: # 10% Keep Original if random.random() < 0.5: masked_token = tokens[index] # 10% replace w/ random word else: masked_token = self.vocab_words[random.randint( 0, len(self.vocab_words) - 1)] output_tokens[index] = masked_token masked_lms.append( MaskedLMInstance(index=index, label=tokens[index])) masked_lms = sorted(masked_lms, key=lambda x: x.index) masked_lm_output = [-1] * len(output_tokens) for p in masked_lms: masked_lm_output[p.index] = self.tokenizer.vocab[p.label] return (output_tokens, masked_lm_output) def get_new_segment(self, segment): """ 输入一句话，返回一句经过处理的话: 为了支持中文全称mask，将被分开的词，将上特殊标记("#")，使得后续处理模块，能够知道哪些字是属于同一个词的。 :param segment: 一句话 :return: 一句处理过的话 """ seq_cws = jieba.lcut(''.join(segment)) seq_cws_dict = {x: 1 for x in seq_cws} new_segment = [] i = 0 while i < len(segment): if len(self.rec.findall(segment[i])) == 0: # 不是中文的，原文加进去。 new_segment.append(segment[i]) i += 1 continue has_add = False for length in range(3, 0, -1): if i + length > len(segment): continue if ''.join(segment[i: i+length]) in seq_cws_dict: new_segment.append(segment[i]) for l in range(1, length): new_segment.append('##' + segment[i+l]) i += length has_add = True break if not has_add: new_segment.append(segment[i]) i += 1 return new_segment def create_whole_masked_lm_predictions(self, tokens): """Creates the predictions for the masked LM objective.""" cand_indexes = [] for (i, token) in enumerate(tokens): if token == "[CLS]" or token == "[SEP]": continue # Whole Word Masking means that if we mask all of the wordpieces # corresponding to an original word. When a word has been split into # WordPieces, the first token does not have any marker and any subsequence # tokens are prefixed with ##. So whenever we see the ## token, we # append it to the previous set of word indexes. # # Note that Whole Word Masking does *not* change the training code # at all -- we still predict each WordPiece independently, softmaxed # over the entire vocabulary. if (self.do_whole_word_mask and len(cand_indexes) >= 1 and token.startswith("##")): cand_indexes[-1].append(i) else: cand_indexes.append([i]) random.shuffle(cand_indexes) output_tokens = [t[2:] if len(self.whole_rec.findall(t))>0 else t for t in tokens] # 去掉"##" num_to_predict = min(self.max_predictions_per_seq, max(1, int(round(len(tokens) * self.masked_lm_prob)))) masked_lms = [] covered_indexes = set() for index_set in cand_indexes: if len(masked_lms) >= num_to_predict: break # If adding a whole-word mask would exceed the maximum number of # predictions, then just skip this candidate. if len(masked_lms) + len(index_set) > num_to_predict: continue is_any_index_covered = False for index in index_set: if index in covered_indexes: is_any_index_covered = True break if is_any_index_covered: continue for index in index_set: covered_indexes.add(index) masked_token = None # 80% of the time, replace with [MASK] if random.random() < 0.8: masked_token = "[MASK]" else: # 10% of the time, keep original if random.random() < 0.5: masked_token = tokens[index][2:] if len(self.whole_rec.findall(tokens[index]))>0 else tokens[index] # 去掉"##" # 10% of the time, replace with random word else: masked_token = self.vocab_words[random.randint(0, len(self.vocab_words) - 1)] output_tokens[index] = masked_token masked_lms.append(MaskedLMInstance(index=index, label=tokens[index][2:] if len(self.whole_rec.findall(tokens[index]))>0 else tokens[index])) assert len(masked_lms) <= num_to_predict masked_lms = sorted(masked_lms, key=lambda x: x.index) masked_lm_output = [-1] * len(output_tokens) for p in masked_lms: masked_lm_output[p.index] = self.tokenizer.vocab[p.label] return (output_tokens, masked_lm_output)