def __init__(self, chpt_path, max_length=128): """ Initializer """ self.checkpoint_path = chpt_path self.bert_pretrained_checkpoint = 'bert-base-uncased' self.max_length = max_length self.sentence_encoder = BERTPAIRSentenceEncoder( self.bert_pretrained_checkpoint, self.max_length) self.model = Pair(self.sentence_encoder, hidden_size=768) if torch.cuda.is_available(): self.model = self.model.cuda() self.model.eval() # self.nlp_coref = spacy.load("en_core_web_sm") # neuralcoref.add_to_pipe(self.nlp_coref) self.nlp_no_coref = spacy.load("en_core_web_sm") self.load_model()
def main(): opt = {'train': 'train_wiki', 'val': 'val_wiki', 'test': 'val_semeval', 'adv': None, 'trainN': 10, 'N': 5, 'K': 1, 'Q': 1, 'batch_size': 1, 'train_iter': 5, 'val_iter': 1000, 'test_iter': 10, 'val_step': 2000, 'model': 'pair', 'encoder': 'bert', 'max_length': 64, 'lr': -1, 'weight_decay': 1e-5, 'dropout': 0.0, 'na_rate':0, 'optim': 'adam', 'load_ckpt': './src/FewRel-master/Checkpoints/ckpt_5-Way-1-Shot_FewRel.pth', 'save_ckpt': './src/FewRel-master/Checkpoints/post_ckpt_5-Way-1-Shot_FewRel.pth', 'fp16':False, 'only_test': False, 'ckpt_name': 'Checkpoints/ckpt_5-Way-1-Shot_FewRel.pth', 'pair': True, 'pretrain_ckpt': '', 'cat_entity_rep': False, 'dot': False, 'no_dropout': False, 'mask_entity': False, 'use_sgd_for_bert': False } opt = DotMap(opt) trainN = opt.trainN N = opt.N K = opt.K Q = opt.Q batch_size = opt.batch_size model_name = opt.model encoder_name = opt.encoder max_length = opt.max_length print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K)) print("model: {}".format(model_name)) print("encoder: {}".format(encoder_name)) print("max_length: {}".format(max_length)) pretrain_ckpt = opt.pretrain_ckpt or 'bert-base-uncased' sentence_encoder = BERTPAIRSentenceEncoder( pretrain_ckpt, max_length) train_data_loader = get_loader_pair(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size, encoder_name=encoder_name) val_data_loader = get_loader_pair(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size, encoder_name=encoder_name) test_data_loader = get_loader_pair(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size, encoder_name=encoder_name) if opt.optim == 'sgd': pytorch_optim = optim.SGD elif opt.optim == 'adam': pytorch_optim = optim.Adam elif opt.optim == 'adamw': pytorch_optim = AdamW else: raise NotImplementedError framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader) prefix = '-'.join([model_name, encoder_name, opt.train, opt.val, str(N), str(K)]) if opt.na_rate != 0: prefix += '-na{}'.format(opt.na_rate) if opt.dot: prefix += '-dot' if opt.cat_entity_rep: prefix += '-catentity' if len(opt.ckpt_name) > 0: prefix += '-' + opt.ckpt_name model = Pair(sentence_encoder, hidden_size=opt.hidden_size) if not os.path.exists('checkpoint'): os.mkdir('checkpoint') ckpt = 'checkpoint/{}.pth.tar'.format(prefix) if opt.save_ckpt: ckpt = opt.save_ckpt if torch.cuda.is_available(): model.cuda() if not opt.only_test: if encoder_name in ['bert', 'roberta']: bert_optim = True else: bert_optim = False if opt.lr == -1: if bert_optim: opt.lr = 2e-5 else: opt.lr = 1e-1 framework.train(model, prefix, batch_size, trainN, N, K, Q, pytorch_optim=pytorch_optim, load_ckpt=opt.load_ckpt, save_ckpt=ckpt, na_rate=opt.na_rate, val_step=opt.val_step, fp16=opt.fp16, pair=opt.pair, train_iter=opt.train_iter, val_iter=opt.val_iter, bert_optim=bert_optim, learning_rate=opt.lr, use_sgd_for_bert=opt.use_sgd_for_bert) else: ckpt = opt.load_ckpt if ckpt is None: print("Warning: --load_ckpt is not specified. Will load Huggingface pre-trained checkpoint.") ckpt = 'none' acc = framework.eval(model, batch_size, N, K, Q, opt.test_iter, na_rate=opt.na_rate, ckpt=ckpt, pair=opt.pair) print("RESULT: %.2f" % (acc * 100))
def main(): parser = argparse.ArgumentParser() parser.add_argument('--train', default='train_wiki', help='train file') parser.add_argument('--val', default='val_wiki', help='val file') parser.add_argument('--test', default='test_wiki', help='test file') parser.add_argument('--adv', default=None, help='adv file') parser.add_argument('--trainN', default=10, type=int, help='N in train') parser.add_argument('--N', default=5, type=int, help='N way') parser.add_argument('--K', default=5, type=int, help='K shot') parser.add_argument('--Q', default=5, type=int, help='Num of query per class') parser.add_argument('--batch_size', default=4, type=int, help='batch size') parser.add_argument('--train_iter', default=30000, type=int, help='num of iters in training') parser.add_argument('--val_iter', default=1000, type=int, help='num of iters in validation') parser.add_argument('--test_iter', default=10000, type=int, help='num of iters in testing') parser.add_argument('--val_step', default=2000, type=int, help='val after training how many iters') parser.add_argument('--model', default='proto', help='model name') parser.add_argument('--encoder', default='cnn', help='encoder: cnn or bert or roberta') parser.add_argument('--max_length', default=128, type=int, help='max length') parser.add_argument('--lr', default=1e-1, type=float, help='learning rate') parser.add_argument('--weight_decay', default=1e-5, type=float, help='weight decay') parser.add_argument('--dropout', default=0.0, type=float, help='dropout rate') parser.add_argument('--na_rate', default=0, type=int, help='NA rate (NA = Q * na_rate)') parser.add_argument('--grad_iter', default=1, type=int, help='accumulate gradient every x iterations') parser.add_argument('--optim', default='sgd', help='sgd / adam / adamw') parser.add_argument('--hidden_size', default=230, type=int, help='hidden size') parser.add_argument('--load_ckpt', default=None, help='load ckpt') parser.add_argument('--save_ckpt', default=None, help='save ckpt') parser.add_argument('--fp16', action='store_true', help='use nvidia apex fp16') parser.add_argument('--only_test', action='store_true', help='only test') # only for bert / roberta parser.add_argument('--pair', action='store_true', help='use pair model') parser.add_argument('--pretrain_ckpt', default=None, help='bert / roberta pre-trained checkpoint') parser.add_argument( '--cat_entity_rep', action='store_true', help='concatenate entity representation as sentence rep') # only for prototypical networks parser.add_argument('--dot', action='store_true', help='use dot instead of L2 distance for proto') opt = parser.parse_args() trainN = opt.trainN N = opt.N K = opt.K Q = opt.Q batch_size = opt.batch_size model_name = opt.model encoder_name = opt.encoder max_length = opt.max_length print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K)) print("model: {}".format(model_name)) print("encoder: {}".format(encoder_name)) print("max_length: {}".format(max_length)) if encoder_name == 'cnn': try: glove_mat = np.load('./pretrain/glove/glove_mat.npy') glove_word2id = json.load( open('./pretrain/glove/glove_word2id.json')) except: raise Exception( "Cannot find glove files. Run glove/download_glove.sh to download glove files." ) sentence_encoder = CNNSentenceEncoder(glove_mat, glove_word2id, max_length) elif encoder_name == 'bert': pretrain_ckpt = opt.pretrain_ckpt or 'bert-base-uncased' if opt.pair: sentence_encoder = BERTPAIRSentenceEncoder(pretrain_ckpt, max_length) else: sentence_encoder = BERTSentenceEncoder( pretrain_ckpt, max_length, cat_entity_rep=opt.cat_entity_rep) elif encoder_name == 'roberta': pretrain_ckpt = opt.pretrain_ckpt or 'roberta-base' if opt.pair: sentence_encoder = RobertaPAIRSentenceEncoder( pretrain_ckpt, max_length) else: sentence_encoder = RobertaSentenceEncoder( pretrain_ckpt, max_length, cat_entity_rep=opt.cat_entity_rep) else: raise NotImplementedError if opt.pair: train_data_loader = get_loader_pair(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size, encoder_name=encoder_name) val_data_loader = get_loader_pair(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size, encoder_name=encoder_name) test_data_loader = get_loader_pair(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size, encoder_name=encoder_name) else: train_data_loader = get_loader(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.adv: adv_data_loader = get_loader_unsupervised(opt.adv, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.optim == 'sgd': pytorch_optim = optim.SGD elif opt.optim == 'adam': pytorch_optim = optim.Adam elif opt.optim == 'adamw': from transformers import AdamW pytorch_optim = AdamW else: raise NotImplementedError if opt.adv: d = Discriminator(opt.hidden_size) framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader, adv_data_loader, adv=opt.adv, d=d) else: framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader) prefix = '-'.join( [model_name, encoder_name, opt.train, opt.val, str(N), str(K)]) if opt.adv is not None: prefix += '-adv_' + opt.adv if opt.na_rate != 0: prefix += '-na{}'.format(opt.na_rate) if opt.dot: prefix += '-dot' if opt.cat_entity_rep: prefix += '-catentity' if model_name == 'proto': model = Proto(sentence_encoder, dot=opt.dot) elif model_name == 'gnn': model = GNN(sentence_encoder, N, hidden_size=opt.hidden_size) elif model_name == 'snail': model = SNAIL(sentence_encoder, N, K, hidden_size=opt.hidden_size) elif model_name == 'metanet': model = MetaNet(N, K, sentence_encoder.embedding, max_length) elif model_name == 'siamese': model = Siamese(sentence_encoder, hidden_size=opt.hidden_size, dropout=opt.dropout) elif model_name == 'pair': model = Pair(sentence_encoder, hidden_size=opt.hidden_size) else: raise NotImplementedError if not os.path.exists('checkpoint'): os.mkdir('checkpoint') ckpt = 'checkpoint/{}.pth.tar'.format(prefix) if opt.save_ckpt: ckpt = opt.save_ckpt if torch.cuda.is_available(): model.cuda() if not opt.only_test: if encoder_name in ['bert', 'roberta']: bert_optim = True else: bert_optim = False framework.train(model, prefix, batch_size, trainN, N, K, Q, pytorch_optim=pytorch_optim, load_ckpt=opt.load_ckpt, save_ckpt=ckpt, na_rate=opt.na_rate, val_step=opt.val_step, fp16=opt.fp16, pair=opt.pair, train_iter=opt.train_iter, val_iter=opt.val_iter, bert_optim=bert_optim) else: ckpt = opt.load_ckpt acc = framework.eval(model, batch_size, N, K, Q, opt.test_iter, na_rate=opt.na_rate, ckpt=ckpt, pair=opt.pair) print("RESULT: %.2f" % (acc * 100))
def main(): parser = argparse.ArgumentParser() parser.add_argument('--train', default='train_wiki', help='train file') parser.add_argument('--val', default='val_wiki', help='val file') parser.add_argument('--test', default='test_wiki', help='test file') parser.add_argument('--adv', default=None, help='adv file') parser.add_argument('--trainN', default=10, type=int, help='N in train') parser.add_argument('--N', default=5, type=int, help='N way') parser.add_argument('--K', default=5, type=int, help='K shot') parser.add_argument('--Q', default=5, type=int, help='Num of query per class') parser.add_argument('--batch_size', default=4, type=int, help='batch size') parser.add_argument('--train_iter', default=20000, type=int, help='num of iters in training') parser.add_argument('--val_iter', default=1000, type=int, help='num of iters in validation') parser.add_argument('--test_iter', default=2000, type=int, help='num of iters in testing') parser.add_argument('--val_step', default=2000, type=int, help='val after training how many iters') parser.add_argument('--model', default='proto', help='model name') parser.add_argument('--encoder', default='cnn', help='encoder: cnn or bert') parser.add_argument('--max_length', default=128, type=int, help='max length') parser.add_argument('--lr', default=1e-1, type=float, help='learning rate') parser.add_argument('--weight_decay', default=1e-5, type=float, help='weight decay') parser.add_argument('--dropout', default=0.0, type=float, help='dropout rate') parser.add_argument('--na_rate', default=0, type=int, help='NA rate (NA = Q * na_rate)') parser.add_argument('--grad_iter', default=1, type=int, help='accumulate gradient every x iterations') parser.add_argument('--optim', default='sgd', help='sgd / adam / bert_adam') parser.add_argument('--hidden_size', default=230, type=int, help='hidden size') parser.add_argument('--load_ckpt', default=None, help='load ckpt') parser.add_argument('--save_ckpt', default=None, help='save ckpt') parser.add_argument('--fp16', action='store_true', help='use nvidia apex fp16') parser.add_argument('--only_test', action='store_true', help='only test') parser.add_argument('--pair', action='store_true', help='use pair model') parser.add_argument('--language', type=str, default='eng', help='language') parser.add_argument('--sup_cost', type=int, default=0, help='use sup classifier') opt = parser.parse_args() trainN = opt.trainN N = opt.N K = opt.K Q = opt.Q batch_size = opt.batch_size model_name = opt.model encoder_name = opt.encoder max_length = opt.max_length sup_cost = bool(opt.sup_cost) print(sup_cost) print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K)) print("model: {}".format(model_name)) print("encoder: {}".format(encoder_name)) print("max_length: {}".format(max_length)) embsize = 50 if opt.language == 'chn': embsize = 100 if encoder_name == 'cnn': try: if opt.language == 'chn': glove_mat = np.load('./pretrain/chinese_emb/emb.npy') glove_word2id = json.load( open('./pretrain/chinese_emb/word2id.json')) else: glove_mat = np.load('./pretrain/glove/glove_mat.npy') glove_word2id = json.load( open('./pretrain/glove/glove_word2id.json')) except: raise Exception( "Cannot find glove files. Run glove/download_glove.sh to download glove files." ) sentence_encoder = CNNSentenceEncoder(glove_mat, glove_word2id, max_length, word_embedding_dim=embsize) elif encoder_name == 'bert': if opt.pair: if opt.language == 'chn': sentence_encoder = BERTPAIRSentenceEncoder( 'bert-base-chinese', #'./pretrain/bert-base-uncased', max_length) else: sentence_encoder = BERTPAIRSentenceEncoder( 'bert-base-uncased', max_length) else: if opt.language == 'chn': sentence_encoder = BERTSentenceEncoder( 'bert-base-chinese', #'./pretrain/bert-base-uncased', max_length) else: sentence_encoder = BERTSentenceEncoder('bert-base-uncased', max_length) else: raise NotImplementedError if opt.pair: train_data_loader = get_loader_pair(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader_pair(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader_pair(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) else: train_data_loader = get_loader(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.adv: adv_data_loader = get_loader_unsupervised(opt.adv, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.optim == 'sgd': pytorch_optim = optim.SGD elif opt.optim == 'adam': pytorch_optim = optim.Adam elif opt.optim == 'bert_adam': from transformers import AdamW pytorch_optim = AdamW else: raise NotImplementedError if opt.adv: d = Discriminator(opt.hidden_size) framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader, adv_data_loader, adv=opt.adv, d=d) else: framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader) prefix = '-'.join( [model_name, encoder_name, opt.train, opt.val, str(N), str(K)]) if opt.adv is not None: prefix += '-adv_' + opt.adv if opt.na_rate != 0: prefix += '-na{}'.format(opt.na_rate) if model_name == 'proto': model = Proto(sentence_encoder, hidden_size=opt.hidden_size) elif model_name == 'gnn': model = GNN(sentence_encoder, N, use_sup_cost=sup_cost) elif model_name == 'snail': print("HINT: SNAIL works only in PyTorch 0.3.1") model = SNAIL(sentence_encoder, N, K) elif model_name == 'metanet': model = MetaNet(N, K, sentence_encoder.embedding, max_length, use_sup_cost=sup_cost) elif model_name == 'siamese': model = Siamese(sentence_encoder, hidden_size=opt.hidden_size, dropout=opt.dropout) elif model_name == 'pair': model = Pair(sentence_encoder, hidden_size=opt.hidden_size) else: raise NotImplementedError if not os.path.exists('checkpoint'): os.mkdir('checkpoint') ckpt = 'checkpoint/{}.pth.tar'.format(prefix) if opt.save_ckpt: ckpt = opt.save_ckpt if torch.cuda.is_available(): model.cuda() if not opt.only_test: if encoder_name == 'bert': bert_optim = True else: bert_optim = False framework.train(model, prefix, batch_size, trainN, N, K, Q, pytorch_optim=pytorch_optim, load_ckpt=opt.load_ckpt, save_ckpt=ckpt, na_rate=opt.na_rate, val_step=opt.val_step, fp16=opt.fp16, pair=opt.pair, train_iter=opt.train_iter, val_iter=opt.val_iter, bert_optim=bert_optim, sup_cls=sup_cost) else: ckpt = opt.load_ckpt acc = framework.eval(model, batch_size, N, K, Q, opt.test_iter, na_rate=opt.na_rate, ckpt=ckpt, pair=opt.pair) wfile = open('logs/' + ckpt.replace('checkpoint/', '') + '.txt', 'a') wfile.write(str(N) + '\t' + str(K) + '\t' + str(acc * 100) + '\n') wfile.close() print("RESULT: %.2f" % (acc * 100))
def main(): parser = argparse.ArgumentParser() parser.add_argument('--train', default='train_wiki', help='train file') parser.add_argument('--val', default='val_wiki', help='val file') parser.add_argument('--test', default='test_wiki', help='test file') parser.add_argument('--adv', default=None, help='adv file') parser.add_argument('--trainN', default=10, type=int, help='N in train') parser.add_argument('--N', default=5, type=int, help='N way') parser.add_argument('--K', default=5, type=int, help='K shot') parser.add_argument('--Q', default=5, type=int, help='Num of query per class') parser.add_argument('--batch_size', default=4, type=int, help='batch size') parser.add_argument('--train_iter', default=30000, type=int, help='num of iters in training') parser.add_argument('--val_iter', default=1000, type=int, help='num of iters in validation') parser.add_argument('--test_iter', default=10000, type=int, help='num of iters in testing') parser.add_argument('--val_step', default=2000, type=int, help='val after training how many iters') parser.add_argument('--model', default='regrab', help='model name') parser.add_argument('--encoder', default='bert', help='encoder: cnn or bert') parser.add_argument('--max_length', default=128, type=int, help='max length') parser.add_argument('--lr', default=1e-1, type=float, help='learning rate') parser.add_argument('--weight_decay', default=1e-5, type=float, help='weight decay') parser.add_argument('--dropout', default=0.0, type=float, help='dropout rate') parser.add_argument('--na_rate', default=0, type=int, help='NA rate (NA = Q * na_rate)') parser.add_argument('--grad_iter', default=1, type=int, help='accumulate gradient every x iterations') parser.add_argument('--optim', default='sgd', help='sgd / adam / bert_adam') parser.add_argument('--hidden_size', default=230, type=int, help='hidden size') parser.add_argument('--load_ckpt', default=None, help='load ckpt') parser.add_argument('--save_ckpt', default=None, help='save ckpt') parser.add_argument('--fp16', action='store_true', help='use nvidia apex fp16') parser.add_argument('--only_test', action='store_true', help='only test') parser.add_argument('--pair', action='store_true', help='use pair model') parser.add_argument('--eps', default=0.1, type=float, help='step size for SG-MCMC') parser.add_argument('--temp', default=10.0, type=float, help='temperature for softmax') parser.add_argument('--step', default=5, type=int, help='steps for SG-MCMC') parser.add_argument('--smp', default=10, type=int, help='samples for SG-MCMC') parser.add_argument('--ratio', default=0.01, type=float, help='decay ratio of step size for SG-MCMC') parser.add_argument('--wtp', default=0.1, type=float, help='weight of the prior term') parser.add_argument('--wtn', default=1.0, type=float, help='weight of the noise term') parser.add_argument('--wtb', default=0.0, type=float, help='weight of the background term') parser.add_argument('--metric', default='dot', help='similarity metric (dot or l2)') parser.add_argument('--seed', default=1234, type=int, help='random seed') opt = parser.parse_args() trainN = opt.trainN N = opt.N K = opt.K Q = opt.Q batch_size = opt.batch_size model_name = opt.model encoder_name = opt.encoder max_length = opt.max_length torch.manual_seed(opt.seed) np.random.seed(opt.seed) random.seed(opt.seed) torch.cuda.manual_seed(opt.seed) print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K)) print("model: {}".format(model_name)) print("encoder: {}".format(encoder_name)) print("max_length: {}".format(max_length)) if encoder_name == 'cnn': try: glove_mat = np.load('./pretrain/glove/glove_mat.npy') glove_word2id = json.load( open('./pretrain/glove/glove_word2id.json')) except: raise Exception( "Cannot find glove files. Run glove/download_glove.sh to download glove files." ) sentence_encoder = CNNSentenceEncoder(glove_mat, glove_word2id, max_length) elif encoder_name == 'bert': if opt.pair: sentence_encoder = BERTPAIRSentenceEncoder( './pretrain/bert-base-uncased', max_length) else: sentence_encoder = BERTSentenceEncoder( './pretrain/bert-base-uncased', max_length) else: raise NotImplementedError if opt.pair: train_data_loader = get_loader_pair(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader_pair(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader_pair(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) else: train_data_loader = get_loader(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.adv: adv_data_loader = get_loader_unsupervised(opt.adv, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.optim == 'sgd': pytorch_optim = optim.SGD elif opt.optim == 'adam': pytorch_optim = optim.Adam elif opt.optim == 'bert_adam': from transformers import AdamW pytorch_optim = AdamW else: raise NotImplementedError if opt.adv: d = Discriminator(opt.hidden_size) framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader, adv_data_loader, adv=opt.adv, d=d) else: framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader) prefix = '-'.join( [model_name, encoder_name, opt.train, opt.val, str(N), str(K)]) if opt.adv is not None: prefix += '-adv_' + opt.adv if opt.na_rate != 0: prefix += '-na{}'.format(opt.na_rate) # Loading relation embeddings. data = pickle.load(open('./data/embeddings.pkl', 'rb')) rellist = data['relations'] relemb = data['embeddings'] array0 = np.zeros((1, relemb.shape[1]), dtype=relemb.dtype) relemb = np.concatenate([array0, relemb], axis=0) rel2id = dict([(rel, k + 1) for k, rel in enumerate(rellist)]) # Loading relation graphs. with open('./data/graph.txt', 'r') as fi: us, vs, ws = [], [], [] for line in fi: items = line.strip().split('\t') us += [rel2id[items[0]]] vs += [rel2id[items[1]]] ws += [float(items[2])] index = torch.LongTensor([us, vs]) value = torch.Tensor(ws) shape = torch.Size([len(rel2id) + 1, len(rel2id) + 1]) reladj = torch.sparse.FloatTensor(index, value, shape).cuda() # End model = REGRAB(sentence_encoder, hidden_size=opt.hidden_size, eps=opt.eps, temp=opt.temp, step=opt.step, smp=opt.smp, ratio=opt.ratio, wtp=opt.wtp, wtn=opt.wtn, wtb=opt.wtb, metric=opt.metric) model.set_relemb(rel2id, relemb) model.set_reladj(reladj) if not os.path.exists('checkpoint'): os.mkdir('checkpoint') ckpt = 'checkpoint/{}.pth.tar'.format(prefix) if opt.save_ckpt: ckpt = opt.save_ckpt if torch.cuda.is_available(): model.cuda() if not opt.only_test: if encoder_name == 'bert': bert_optim = True else: bert_optim = False framework.train(model, prefix, batch_size, trainN, N, K, Q, pytorch_optim=pytorch_optim, load_ckpt=opt.load_ckpt, save_ckpt=ckpt, na_rate=opt.na_rate, val_step=opt.val_step, fp16=opt.fp16, pair=opt.pair, train_iter=opt.train_iter, val_iter=opt.val_iter, bert_optim=bert_optim) else: ckpt = opt.load_ckpt acc = framework.eval(model, batch_size, N, K, Q, opt.test_iter, na_rate=opt.na_rate, ckpt=ckpt, pair=opt.pair) print("RESULT: %.2f" % (acc * 100))
def main(): parser = argparse.ArgumentParser() parser.add_argument('--train', default='train_wiki', help='train file') parser.add_argument('--val', default='val_wiki', help='val file') parser.add_argument('--test', default='test_wiki', help='test file') parser.add_argument('--adv', default=None, help='adv file') parser.add_argument('--trainN', default=10, type=int, help='N in train') parser.add_argument('--N', default=5, type=int, help='N way') parser.add_argument('--K', default=5, type=int, help='K shot') parser.add_argument('--Q', default=5, type=int, help='Num of query per class') parser.add_argument('--batch_size', default=4, type=int, help='batch size') parser.add_argument('--train_iter', default=30000, type=int, help='num of iters in training') parser.add_argument('--val_iter', default=1000, type=int, help='num of iters in validation') parser.add_argument('--test_iter', default=3000, type=int, help='num of iters in testing') parser.add_argument('--val_step', default=2000, type=int, help='val after training how many iters') parser.add_argument('--model', default='proto', help='model name') parser.add_argument('--encoder', default='cnn', help='encoder: cnn or bert') parser.add_argument('--max_length', default=128, type=int, help='max length') parser.add_argument('--lr', default=1e-1, type=float, help='learning rate') parser.add_argument('--weight_decay', default=1e-5, type=float, help='weight decay') parser.add_argument('--dropout', default=0.0, type=float, help='dropout rate') parser.add_argument('--na_rate', default=0, type=int, help='NA rate (NA = Q * na_rate)') parser.add_argument('--grad_iter', default=1, type=int, help='accumulate gradient every x iterations') parser.add_argument('--optim', default='sgd', help='sgd / adam / bert_adam') parser.add_argument('--hidden_size', default=230, type=int, help='hidden size') parser.add_argument('--load_ckpt', default=None, help='load ckpt') parser.add_argument('--save_ckpt', default=None, help='save ckpt') parser.add_argument('--fp16', action='store_true', help='use nvidia apex fp16') parser.add_argument('--only_test', action='store_true', help='only test') parser.add_argument('--pair', action='store_true', help='use pair model') opt = parser.parse_args() trainN = opt.trainN N = opt.N K = opt.K Q = opt.Q batch_size = opt.batch_size model_name = opt.model encoder_name = opt.encoder max_length = opt.max_length print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K)) print("model: {}".format(model_name)) print("encoder: {}".format(encoder_name)) print("max_length: {}".format(max_length)) if encoder_name == 'cnn': try: glove_mat = np.load('./pretrain/glove/glove_mat.npy') glove_word2id = json.load( open('./pretrain/glove/glove_word2id.json')) except: raise Exception( "Cannot find glove files. Run glove/download_glove.sh to download glove files." ) sentence_encoder = CNNSentenceEncoder(glove_mat, glove_word2id, max_length) elif encoder_name == 'bert': if opt.pair: sentence_encoder = BERTPAIRSentenceEncoder( './pretrain/bert-base-uncased', max_length) else: sentence_encoder = BERTSentenceEncoder( './pretrain/bert-base-uncased', max_length) else: raise NotImplementedError if opt.pair: train_data_loader = get_loader_pair(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader_pair(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader_pair(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) else: train_data_loader = get_loader(opt.train, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) val_data_loader = get_loader(opt.val, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) test_data_loader = get_loader(opt.test, sentence_encoder, N=N, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.adv: adv_data_loader = get_loader_unsupervised(opt.adv, sentence_encoder, N=trainN, K=K, Q=Q, na_rate=opt.na_rate, batch_size=batch_size) if opt.optim == 'sgd': pytorch_optim = optim.SGD elif opt.optim == 'adam': pytorch_optim = optim.Adam elif opt.optim == 'bert_adam': from pytorch_transformers import AdamW pytorch_optim = AdamW else: raise NotImplementedError if opt.adv: d = Discriminator(opt.hidden_size) framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader, adv_data_loader, adv=opt.adv, d=d) else: framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader) prefix = '-'.join( [model_name, encoder_name, opt.train, opt.val, str(N), str(K)]) if opt.adv is not None: prefix += '-adv_' + opt.adv if opt.na_rate != 0: prefix += '-na{}'.format(opt.na_rate) if model_name == 'proto': model = Proto(sentence_encoder, hidden_size=opt.hidden_size) elif model_name == 'gnn': model = GNN(sentence_encoder, N) elif model_name == 'snail': print("HINT: SNAIL works only in PyTorch 0.3.1") model = SNAIL(sentence_encoder, N, K) elif model_name == 'metanet': model = MetaNet(N, K, sentence_encoder.embedding, max_length) elif model_name == 'siamese': model = Siamese(sentence_encoder, hidden_size=opt.hidden_size, dropout=opt.dropout) elif model_name == 'pair': model = Pair(sentence_encoder, hidden_size=opt.hidden_size) else: raise NotImplementedError if not os.path.exists('checkpoint'): os.mkdir('checkpoint') ckpt = 'checkpoint/{}.pth.tar'.format(prefix) if opt.save_ckpt: ckpt = opt.save_ckpt if torch.cuda.is_available(): model.cuda() if not opt.only_test: if encoder_name == 'bert': bert_optim = True else: bert_optim = False framework.train(model, prefix, batch_size, trainN, N, K, Q, pytorch_optim=pytorch_optim, load_ckpt=opt.load_ckpt, save_ckpt=ckpt, na_rate=opt.na_rate, val_step=opt.val_step, fp16=opt.fp16, pair=opt.pair, train_iter=opt.train_iter, val_iter=opt.val_iter, bert_optim=bert_optim) else: ckpt = opt.load_ckpt acc = 0 his_acc = [] total_test_round = 5 for i in range(total_test_round): cur_acc = framework.eval(model, batch_size, N, K, Q, opt.test_iter, na_rate=opt.na_rate, ckpt=ckpt, pair=opt.pair) his_acc.append(cur_acc) acc += cur_acc acc /= total_test_round nhis_acc = np.array(his_acc) error = nhis_acc.std() * 1.96 / (nhis_acc.shape[0]**0.5) print("RESULT: %.2f\\pm%.2f" % (acc * 100, error * 100)) result_file = open('./result.txt', 'a+') result_file.write( "test data: %12s | model: %45s | acc: %.6f\n | error: %.6f\n" % (opt.test, prefix, acc, error)) result_file = open('./result_detail.txt', 'a+') result_detail = { 'test': opt.test, 'model': prefix, 'acc': acc, 'his': his_acc } result_file.write("%s\n" % (json.dumps(result_detail)))
class Detector: #RUNS USING GPU if available and pytoch has CUDA support def __init__(self, chpt_path, max_length=128): """ Initializer """ self.checkpoint_path = chpt_path self.bert_pretrained_checkpoint = 'bert-base-uncased' self.max_length = max_length self.sentence_encoder = BERTPAIRSentenceEncoder( self.bert_pretrained_checkpoint, self.max_length) self.model = Pair(self.sentence_encoder, hidden_size=768) if torch.cuda.is_available(): self.model = self.model.cuda() self.model.eval() # self.nlp_coref = spacy.load("en_core_web_sm") # neuralcoref.add_to_pipe(self.nlp_coref) self.nlp_no_coref = spacy.load("en_core_web_sm") self.load_model() def __load_model_from_checkpoint__(self, ckpt): ''' ckpt: Path of the checkpoint return: Checkpoint dict ''' if os.path.isfile(ckpt): checkpoint = torch.load(ckpt) print("Successfully loaded checkpoint '%s'" % ckpt) return checkpoint else: raise Exception("No checkpoint found at '%s'" % ckpt) def bert_tokenize(self, tokens, head_indices, tail_indices): word = self.sentence_encoder.tokenize(tokens, head_indices, tail_indices) return word def load_model(self): """ Loads the model from the checkpoint """ state_dict = self.__load_model_from_checkpoint__( self.checkpoint_path)['state_dict'] own_state = self.model.state_dict() for name, param in state_dict.items(): if name not in own_state: continue own_state[name].copy_(param) # def spacy_tokenize_coref(self,sentence): # """ # Tokenizes the sentence using spacy # """ # return list(map(str, self.nlp_coref(sentence))) def spacy_tokenize_no_coref(self, sentence): """ Tokenizes the sentence using spacy """ try: return list(map(str, self.nlp_no_coref(sentence))) except TypeError as e: print("problem sentence: '{}'".format(sentence)) raise e # def get_head_tail_pairs(self,sentence): # """ # Gets pairs of heads and tails of named entities so that relation identification can be done on these. # """ # acceptable_entity_types = ['PERSON', 'NORP', 'ORG', 'GPE', 'PRODUCT', 'EVENT', 'LAW', 'LOC', 'FAC'] # doc = self.nlp_coref(sentence) # entity_info = [(X.text, X.label_) for X in doc.ents] # entity_info = set(map(lambda x:x[0], filter(lambda x:x[1] in acceptable_entity_types, entity_info))) # return combinations(entity_info, 2) def _get_indices_alt(self, tokens, tokenized_head, tokenized_tail): """ Alternative implemention for getting the indices of the head and tail if exact matches cannot be done. """ head_indices = None tail_indices = None for i in range(len(tokens)): if tokens[i] in tokenized_head[0] or tokenized_head[0] in tokens[i]: broke = False for k, j in zip(tokens[i:i + len(tokenized_head)], tokenized_head): if k not in j and j not in k: broke = True break if not broke: head_indices = list(range(i, i + len(tokenized_head))) break for i in range(len(tokens)): if tokens[i] in tokenized_tail[0] or tokenized_tail[0] in tokens[i]: broke = False for k, j in zip(tokens[i:i + len(tokenized_tail)], tokenized_tail): if k not in j and j not in k: broke = True break if not broke: tail_indices = list(range(i, i + len(tokenized_tail))) break return head_indices, tail_indices def _calculate_conf(self, logits, order, pred): exp = list(float(i) for i in logits[0][0]) exp = [math.exp(i) for i in exp] if pred == 'NA': return exp[-1] * 100 / sum(exp) return exp[order.index(pred)] * 100 / sum(exp) def run_detection_algorithm(self, query, relation_data): """ Runs the algorithm/model on the given query using the given support data. """ N = len(relation_data) K = len(relation_data[0]['examples']) Q = 1 head = query['head'] tail = query['tail'] fusion_set = {'word': [], 'mask': [], 'seg': []} tokens = self.spacy_tokenize_no_coref(query['sentence']) print("head: '{}' tail: '{}' sentence: '{}'".format( head, tail, query['sentence'])) tokenized_head = self.spacy_tokenize_no_coref(head) tokenized_tail = self.spacy_tokenize_no_coref(tail) head_indices = None tail_indices = None for i in range(len(tokens)): if tokens[i] == tokenized_head[0] and tokens[ i:i + len(tokenized_head)] == tokenized_head: head_indices = list(range(i, i + len(tokenized_head))) break for i in range(len(tokens)): if tokens[i] == tokenized_tail[0] and tokens[ i:i + len(tokenized_tail)] == tokenized_tail: tail_indices = list(range(i, i + len(tokenized_tail))) break if head_indices is None or tail_indices is None: head_indices, tail_indices = self._get_indices_alt( tokens, tokenized_head, tokenized_tail) if head_indices is None or tail_indices is None: print(tokenized_head) print(tokenized_tail) print(tokens) raise ValueError( "Head/Tail indices error: head: {} \n tail: {} \n sentence: {}" .format(head, tail, query['sentence'])) bert_query_tokens = self.bert_tokenize(tokens, head_indices, tail_indices) for relation in relation_data: for ex in relation['examples']: tokens = self.spacy_tokenize_no_coref(ex['sentence']) tokenized_head = self.spacy_tokenize_no_coref( ex['head'] ) #head and tail spelling and punctuation should match the corefered output exactly tokenized_tail = self.spacy_tokenize_no_coref(ex['tail']) head_indices = None tail_indices = None for i in range(len(tokens)): if tokens[i] == tokenized_head[0] and tokens[ i:i + len(tokenized_head)] == tokenized_head: head_indices = list(range(i, i + len(tokenized_head))) break for i in range(len(tokens)): if tokens[i] == tokenized_tail[0] and tokens[ i:i + len(tokenized_tail)] == tokenized_tail: tail_indices = list(range(i, i + len(tokenized_tail))) break if head_indices is None or tail_indices is None: raise ValueError( "Head/Tail indices error: head: {} \n tail: {} \n sentence: {}" .format(ex['head'], ex['tail'], ex['sentence'])) bert_relation_example_tokens = self.bert_tokenize( tokens, head_indices, tail_indices) SEP = self.sentence_encoder.tokenizer.convert_tokens_to_ids( ['[SEP]']) CLS = self.sentence_encoder.tokenizer.convert_tokens_to_ids( ['[CLS]']) word_tensor = torch.zeros((self.max_length)).long() new_word = CLS + bert_relation_example_tokens + SEP + bert_query_tokens + SEP for i in range(min(self.max_length, len(new_word))): word_tensor[i] = new_word[i] mask_tensor = torch.zeros((self.max_length)).long() mask_tensor[:min(self.max_length, len(new_word))] = 1 seg_tensor = torch.ones((self.max_length)).long() seg_tensor[:min(self.max_length, len(bert_relation_example_tokens) + 1)] = 0 fusion_set['word'].append(word_tensor) fusion_set['mask'].append(mask_tensor) fusion_set['seg'].append(seg_tensor) fusion_set['word'] = torch.stack(fusion_set['word']) fusion_set['seg'] = torch.stack(fusion_set['seg']) fusion_set['mask'] = torch.stack(fusion_set['mask']) if torch.cuda.is_available(): fusion_set['word'] = fusion_set['word'].cuda() fusion_set['seg'] = fusion_set['seg'].cuda() fusion_set['mask'] = fusion_set['mask'].cuda() logits, pred = self.model(fusion_set, N, K, Q) gc.collect() order = list(r['name'] for r in relation_data) pred_relation = relation_data[ pred.item()]['name'] if pred.item() < len(relation_data) else 'NA' return { 'sentence': query['sentence'], 'head': head, 'tail': tail, 'pred_relation': pred_relation, 'conf': int(self._calculate_conf(logits, order, pred_relation)) } #returns (sentence, head, tail, prediction relation name) def print_result(self, sentence, head, tail, prediction): """ Helper function to print the results to the stdout. """ print('Sentence: \"{}\", head: \"{}\", tail: \"{}\", prediction: {}'. format(sentence, head, tail, prediction))