CCKS比赛.zip

import pickle from NER_MODEL import * from transformers import BertTokenizer, RobertaForTokenClassification tokenizer = BertTokenizer.from_pretrained("RoBERTa_zh_Large_PyTorch") max_seq_len = 40 train_corpus = pickle.load(open('./data/corpus_train.pkl', 'rb')) train_questions = [train_corpus[i]['question'] for i in range(len(train_corpus))] train_entitys = [train_corpus[i]['gold_entitys'] for i in range(len(train_corpus))] train_entitys = [[entity[1:-1].split('_')[0] for entity in line] for line in train_entitys] test_corpus = pickle.load(open('./data/corpus_test.pkl', 'rb')) test_questions = [test_corpus[i]['question'] for i in range(len(test_corpus))] test_entitys = [test_corpus[i]['gold_entitys'] for i in range(len(test_corpus))] test_entitys = [[entity[1:-1].split('_')[0] for entity in line] for line in test_entitys] def find_lcsubstr(s1, s2): m = [[0 for i in range(len(s2) + 1)] for j in range(len(s1) + 1)] # 生成0矩阵，为方便后续计算，比字符串长度多了一列 mmax = 0 # 最长匹配的长度 p = 0 # 最长匹配对应在s1中的最后一位 for i in range(len(s1)): for j in range(len(s2)): if s1[i] == s2[j]: m[i + 1][j + 1] = m[i][j] + 1 if m[i + 1][j + 1] > mmax: mmax = m[i + 1][j + 1] p = i + 1 return s1[p - mmax:p] def GetXY(questions, entitys): X1, X2,X3, Y = [], [], [], [] for i in range(len(questions)): q = questions[i] #x1, x2 = tokenizer.encode(first=q, max_len=max_seq_len) # 分别是 词索引序列和分块索引序列 encoded_dict = tokenizer(q, max_length =max_seq_len,pad_to_max_length=True, return_tensors='pt') # 分别是 词索引序列和分块索引序列 x1, x2, x3 = encoded_dict["input_ids"][0],encoded_dict["token_type_ids"][0],encoded_dict["attention_mask"][0] y = [[0] for j in range(max_seq_len)] assert len(x1) == len(y) for e in entitys[i]: # 得到实体名和问题的最长连续公共子串 e = find_lcsubstr(e, q) if e in q: begin = q.index(e) + 1 end = begin + len(e) if end < max_seq_len - 1: for pos in range(begin, end): y[pos] = [1] X1.append(x1.tolist()) X2.append(x2.tolist()) X3.append(x3.tolist()) Y.append(y) X1 = torch.tensor(X1).long() X2 = torch.tensor(X2).long() X3 = torch.tensor(X3).long() Y = torch.tensor(np.array(Y)).squeeze().long() return X1, X2, X3, Y trainx1, trainx2, trainx3,trainy = GetXY(train_questions, train_entitys) # (num_sample,max_len) testx1, testx2,testx3, testy = GetXY(test_questions, test_entitys) import torch from torch.utils.data import TensorDataset, random_split # 把input 放入 TensorDataset。 train_dataset = TensorDataset(trainx1, trainx2, trainx3,trainy) test_dataset = TensorDataset(testx1, testx2, testx3,testy) from torch.utils.data import DataLoader, RandomSampler, SequentialSampler # 推荐batch_size 为 16 或者 32 batch_size = 8 # 为训练数据集和验证数据集设计DataLoaders. train_dataloader = DataLoader( train_dataset, # 训练数据. sampler = RandomSampler(train_dataset), # 打乱顺序 batch_size = batch_size ) validation_dataloader = DataLoader( test_dataset, # 验证数据. sampler = RandomSampler(test_dataset), # 打乱顺序 batch_size = batch_size ) model = enity_identifing(tokenizer.vocab_size,1000,bert_model) # AdamW 是一个 huggingface library 的类，'W' 是'Weight Decay fix"的意思。 optimizer = AdamW(model.parameters(), lr = 2e-5, # args.learning_rate - 默认是 5e-5 eps = 1e-8 # args.adam_epsilon - 默认是 1e-8， 是为了防止衰减率分母除到0 ) from transformers import get_linear_schedule_with_warmup # bert 推荐 epochs 在2到4之间为好。 epochs = 50 # training steps 的数量: [number of batches] x [number of epochs]. total_steps = len(train_dataloader) * epochs # 设计 learning rate scheduler. scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps = 0, # Default value in run_glue.py num_training_steps = total_steps) def flat_accuracy(preds, labels,attention): scores = (preds*attention == labels*attention) rights = 0 for score in scores: if sum(score) == len(labels[0]): rights += 1 #return np.sum((pred_flat == labels_flat)*atten)/ np.sum(atten) return rights/len(labels) import time import datetime def format_time(elapsed): elapsed_rounded = int(round((elapsed))) # 返回 hh:mm:ss 形式的时间 return str(datetime.timedelta(seconds=elapsed_rounded)) import os import random import numpy as np from transformers import WEIGHTS_NAME, CONFIG_NAME output_dir = './ner1/' output_model_file = os.path.join(output_dir, WEIGHTS_NAME) # 代码参考 https://github.com/huggingface/transformers/blob/5bfcd0485ece086ebcbed2d008813037968a9e58/examples/run_glue.py#L128 # 设置随机种子. seed_val = 42 random.seed(seed_val) np.random.seed(seed_val) torch.manual_seed(seed_val) torch.cuda.manual_seed_all(seed_val) # 记录training ,validation loss ,validation accuracy and timings. maxf = 0.0 def computeF(gold_entity, pre_entity): ''' 根据标注的实体位置和预测的实体位置，计算prf,完全匹配 输入： Python-list 3D，值为每个实体的起始位置列表[begin，end] 输出： float ''' truenum = 0 prenum = 0 goldnum = 0 for i in range(len(gold_entity)): goldnum += len(gold_entity[i]) prenum += len(pre_entity[i]) truenum += len(set(gold_entity[i]).intersection(set(pre_entity[i]))) try: precise = float(truenum) / float(prenum) recall = float(truenum) / float(goldnum) f = float(2 * precise * recall / (precise + recall)) except: precise = recall = f = 0.0 print('本轮实体的F值是 %f' % (f)) return precise, recall, f def restore_entity_from_labels_on_corpus(predicty, questions): def restore_entity_from_labels(labels, question): entitys = [] str = '' labels = labels[1:-1] for i in range(min(len(labels), len(question))): if labels[i] == 1: str += question[i] else: if len(str): entitys.append(str) str = '' if len(str): entitys.append(str) return entitys all_entitys = [] for i in range(len(predicty)): all_entitys.append(restore_entity_from_labels(predicty[i], questions[i])) return all_entitys def train(): training_stats = [] # 设置总时间. total_t0 = time.time() best_val_accuracy = 0 for epoch_i in range(0, epochs): print('Epoch {:} / {:}'.format(epoch_i + 1, epochs)) # 记录每个 epoch 所用的时间 t0 = time.time() total_train_loss = 0 total_train_accuracy = 0 model.train() for step, batch in enumerate(train_dataloader): # 每隔40个batch 输出一下所用时间. if step % 100 == 0 and not step == 0: elapsed = format_time(time.time() - t0) print(' Batch {:>5,} of {:>5,}. Elapsed: {:}.'.format(step, len(train_dataloader), elapsed)) # `batch` 包括3个 tensors: # [0]: input ids # [1]: attention masks # [2]: labels b_input_ids