import os,sys,argparse import pytorch_pretrained_bert from pytorch_pretrained_bert.modeling import BertPreTrainedModel, BertModel, BertConfig from pytorch_pretrained_bert import BertTokenizer from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE import torch from torch import nn from torch.nn import CrossEntropyLoss import torch.optim as optim import sequence_layered_reader, sequence_eval, sequence_reader import numpy as np from tagger import Tagger import flat_reader, layered_reader batch_size=32 dropout_rate=0.25 bert_dim=768 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print('') print("********************************************") print("Running on: {}".format(device)) print("********************************************") print('') def predict(predictFile, outFile, model, ignoreEvents, ignoreEntities): testSentences, origTestSentences = sequence_layered_reader.read_booknlp(predictFile, model) with open(outFile, "w", encoding="utf-8") as out: # FLAT if not ignoreEvents: batched_sents_flat, batched_orig_token_lens_flat, batched_data_flat, batched_mask_flat, batched_transforms_flat, ordering_flat=flat_reader.get_batches(model, testSentences, batch_size, training=False) flat_preds_in_order=model.tagFlat(batched_sents_flat, batched_data_flat, batched_mask_flat, batched_transforms_flat, batched_orig_token_lens_flat, ordering_flat) for s_idx, preds in enumerate(flat_preds_in_order): for w_idx, val in enumerate(preds): (word, label)=val # account for [CLS] o_tid, o_word=origTestSentences[s_idx][w_idx+1] assert (word == o_word) if label == 1: out.write("%s\t%s\t%s\t%s\n" % (o_tid, o_tid, "EVENT", word)) # LAYERED if not ignoreEntities: batched_sents, batched_data, batched_mask, batched_transforms, batched_orig_token_lens, ordering = layered_reader.get_batches(model, testSentences, batch_size, tagset, training=False) preds_in_order=model.tag(batched_sents, batched_data, batched_mask, batched_transforms, batched_orig_token_lens, ordering) for idx, preds in enumerate(preds_in_order): for _, label, start, end in preds: # account for [CLS] start+=1 end+=1 start_token=origTestSentences[idx][start][0] end_token=origTestSentences[idx][end][0] phrase=' '.join([x[0] for x in testSentences[idx][start:end]]) phraseEndToken=int(end_token)-1 if phraseEndToken == -2: phraseEndToken=start_token out.write("%s\t%s\t%s\t%s\n" % (start_token, phraseEndToken, label, phrase)) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--mode', help='{train,test,predict,predictBatch}', required=True) parser.add_argument('--batch_prediction_file', help='Filename containing input paths to tag, paired with output paths to write to', required=False) parser.add_argument('-i','--input_prediction_file', help='Filename to tag', required=False) parser.add_argument('-o','--output_prediction_file', help='Filename to write tagged text to', required=False) parser.add_argument('--trainFolder_flat', help='Folder containing training data (flat)', required=False) parser.add_argument('--testFolder_flat', help='Folder containing test data (flat)', required=False) parser.add_argument('--devFolder_flat', help='Folder containing dev data (flat)', required=False) parser.add_argument('--trainFolder_layered', help='Folder containing training data (layered)', required=False) parser.add_argument('--testFolder_layered', help='Folder containing test data (layered)', required=False) parser.add_argument('--devFolder_layered', help='Folder containing dev data (layered)', required=False) parser.add_argument('--tagFile_flat', help='File mapping tags to tag ids (flat)', default="files/event.tagset", required=False) parser.add_argument('--tagFile_layered', help='File mapping tags to tag ids (layered)', default="files/entity.tagset", required=False) parser.add_argument('--modelFile', help='File to write model to/read from', default="files/event.entity.model", required=False) parser.add_argument('--flat_metric', help='{accuracy,fscore,span_fscore}', required=False) parser.add_argument('--ignoreEvents', action="store_true", default=False, help='tag events', required=False) parser.add_argument('--ignoreEntities', action="store_true", default=False, help='tag entities', required=False) args = vars(parser.parse_args()) print(args) mode=args["mode"] tagset_flat=sequence_layered_reader.read_tagset(args["tagFile_flat"]) tagset=sequence_layered_reader.read_tagset(args["tagFile_layered"]) model_file=args["modelFile"] cache_dir = os.path.join(str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(0)) model = Tagger.from_pretrained('bert-base-cased', cache_dir=cache_dir, freeze_bert=True, tagset_flat=tagset_flat, tagset=tagset, device=device) model.to(device) if mode == "train": train_folder_flat=args["trainFolder_flat"] dev_folder_flat=args["devFolder_flat"] train_folder_layered=args["trainFolder_layered"] dev_folder_layered=args["devFolder_layered"] flat_metric=None if args["flat_metric"].lower() == "fscore": flat_metric=sequence_eval.check_f1_two_lists elif args["flat_metric"].lower() == "accuracy": flat_metric=sequence_eval.get_accuracy elif args["flat_metric"].lower() == "span_fscore": flat_metric=sequence_eval.check_span_f1_two_lists ## Flat trainSentences = sequence_reader.prepare_annotations_from_folder(train_folder_flat, tagset_flat) devSentences = sequence_reader.prepare_annotations_from_folder(dev_folder_flat, tagset_flat) batched_sents_flat, batched_orig_token_lens_flat, batched_data_flat, batched_mask_flat, batched_labels_flat, batched_transforms_flat, ordering_flat=flat_reader.get_batches(model, trainSentences, batch_size) dev_batched_sents_flat, dev_batched_orig_token_lens_flat, dev_batched_data_flat, dev_batched_mask_flat, dev_batched_labels_flat, dev_batched_transforms_flat, dev_ordering_flat=flat_reader.get_batches(model, devSentences, batch_size) # NESTED NER trainSentences_nested = sequence_layered_reader.prepare_annotations_from_folder(train_folder_layered, tagset) devSentences_nested = sequence_layered_reader.prepare_annotations_from_folder(dev_folder_layered, tagset) batched_sents, batched_data, batched_mask, batched_labels, batched_transforms, ordering, batched_layered_labels1, batched_layered_labels2, batched_layered_labels3, batched_layered_labels4, batched_layered_labels5, batched_index1, batched_index2, batched_index3, batched_newlabel1, batched_newlabel2, batched_newlabel3, lens=layered_reader.get_batches(model, trainSentences_nested, batch_size, tagset) dev_batched_sents, dev_batched_data, dev_batched_mask, dev_batched_labels, dev_batched_transforms, dev_ordering, dev_batched_layered_labels1, dev_batched_layered_labels2, dev_batched_layered_labels3, dev_batched_layered_labels4, dev_batched_layered_labels5, dev_batched_index1, dev_batched_index2, dev_batched_index3, dev_batched_newlabel1, dev_batched_newlabel2, dev_batched_newlabel3, dev_lens=layered_reader.get_batches(model, devSentences_nested, batch_size, tagset) optimizer = optim.Adam(model.parameters()) maxScore=0 best_idx=0 patience=10 for epoch in range(100): model.train() for b in range(len(batched_data_flat)): if b % 10 == 0: print(b) sys.stdout.flush() loss_flat = model.forwardFlatSequence(batched_data_flat[b], token_type_ids=None, attention_mask=batched_mask_flat[b], transforms=batched_transforms_flat[b], labels=batched_labels_flat[b]) loss_flat.backward() optimizer.step() model.zero_grad() print("***FLAT EVAL***") score=model.evaluateFlat(dev_batched_data_flat, dev_batched_mask_flat, dev_batched_labels_flat, dev_batched_transforms_flat, flat_metric, tagset_flat) sys.stdout.flush() if score > maxScore: torch.save(model.state_dict(), model_file) print("%s is better than %s, saving..." % (score, maxScore)) maxScore=score best_idx=epoch print(epoch-best_idx, patience, epoch, best_idx) if epoch-best_idx > patience: print ("Stopping flat training at epoch %s" % epoch) break maxScore=0 best_idx=0 patience/=2 # load best event model so far (not current state) model = Tagger.from_pretrained('bert-base-cased', cache_dir=cache_dir, freeze_bert=True, tagset_flat=tagset_flat, tagset=tagset, device=device) model.to(device) model.load_state_dict(torch.load(model_file, map_location=device)) optimizer = optim.Adam(model.parameters()) for epoch in range(100): model.train() for b in range(len(batched_data)): if b % 10 == 0: print(b) sys.stdout.flush() loss_layered = model.forward(batched_data[b], batched_index1[b], batched_index2[b], attention_mask=batched_mask[b], transforms=batched_transforms[b], labels=[batched_newlabel1[b], batched_newlabel2[b], batched_newlabel3[b]], lens=[lens[0][b], lens[1][b], lens[2][b] ]) loss_layered.backward() optimizer.step() model.zero_grad() print("***LAYERED EVAL***") score=model.evaluate(dev_batched_sents, dev_batched_data, dev_batched_mask, dev_batched_labels, dev_batched_transforms, dev_batched_layered_labels1, dev_batched_layered_labels2, dev_batched_layered_labels3, dev_batched_layered_labels4, dev_lens) sys.stdout.flush() if score > maxScore: torch.save(model.state_dict(), model_file) print("%s is better than %s, saving..." % (score, maxScore)) maxScore=score best_idx=epoch print(epoch-best_idx, patience, epoch, best_idx) if epoch-best_idx > patience: print ("Stopping layered training at epoch %s" % epoch) break elif mode == "test": dev_folder_flat=args["testFolder_flat"] dev_folder_layered=args["testFolder_layered"] flat_metric=None if args["flat_metric"].lower() == "fscore": flat_metric=sequence_eval.check_f1_two_lists elif args["flat_metric"].lower() == "accuracy": flat_metric=sequence_eval.get_accuracy elif args["flat_metric"].lower() == "span_fscore": flat_metric=sequence_eval.check_span_f1_two_lists ## FLAT devSentences = sequence_reader.prepare_annotations_from_folder(dev_folder_flat, tagset_flat) dev_batched_sents_flat, dev_batched_orig_token_lens_flat, dev_batched_data_flat, dev_batched_mask_flat, dev_batched_labels_flat, dev_batched_transforms_flat, dev_ordering_flat=flat_reader.get_batches(model, devSentences, batch_size) # NESTED NER devSentences_nested = sequence_layered_reader.prepare_annotations_from_folder(dev_folder_layered, tagset) dev_batched_sents, dev_batched_data, dev_batched_mask, dev_batched_labels, dev_batched_transforms, dev_ordering, dev_batched_layered_labels1, dev_batched_layered_labels2, dev_batched_layered_labels3, dev_batched_layered_labels4, dev_batched_layered_labels5, dev_batched_index1, dev_batched_index2, dev_batched_index3, dev_batched_newlabel1, dev_batched_newlabel2, dev_batched_newlabel3, dev_lens=layered_reader.get_batches(model, devSentences_nested, batch_size, tagset) model.load_state_dict(torch.load(model_file, map_location=device)) print("***FLAT EVAL***") score=model.evaluateFlat(dev_batched_data_flat, dev_batched_mask_flat, dev_batched_labels_flat, dev_batched_transforms_flat, flat_metric, tagset_flat) print("***LAYERED EVAL***") score=model.evaluate(dev_batched_sents, dev_batched_data, dev_batched_mask, dev_batched_labels, dev_batched_transforms, dev_batched_layered_labels1, dev_batched_layered_labels2, dev_batched_layered_labels3, dev_batched_layered_labels4, dev_lens) elif mode == "predict": predictFile=args["input_prediction_file"] outFile=args["output_prediction_file"] model.load_state_dict(torch.load(model_file, map_location=device)) predict(predictFile, outFile, model, bool(args["ignoreEvents"]), bool(args["ignoreEntities"])) elif mode == "predictBatch": pathFile=args["batch_prediction_file"] model.load_state_dict(torch.load(model_file)) inpaths, outpaths=sequence_layered_reader.read_filenames(pathFile) for infile, outfile in zip(inpaths, outpaths): print("Tagging %s" % infile) sys.stdout.flush() predict(infile, outfile, model, bool(args["ignoreEvents"]), bool(args["ignoreEntities"]))