예제 #1
0
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))
if len(sys.argv) > 3:
    K = int(sys.argv[3])

print("{}-way-{}-shot Few-Shot Relation Classification".format(N, K))
print("Model: {}".format(model_name))

max_length = 40
train_data_loader = JSONFileDataLoader('./data/train.json', './data/glove.6B.50d.json', max_length=max_length)
val_data_loader = JSONFileDataLoader('./data/val.json', './data/glove.6B.50d.json', max_length=max_length)
test_data_loader = JSONFileDataLoader('./data/test.json', './data/glove.6B.50d.json', max_length=max_length)

framework = FewShotREFramework(train_data_loader, val_data_loader, test_data_loader)
sentence_encoder = CNNSentenceEncoder(train_data_loader.word_vec_mat, max_length)

if model_name == 'proto':
    model = Proto(sentence_encoder)
    framework.train(model, model_name, 4, 20, N, K, 5)
elif model_name == 'gnn':
    model = GNN(sentence_encoder, N)
    framework.train(model, model_name, 2, N, N, K, 1, learning_rate=1e-3, weight_decay=0, optimizer=optim.Adam)
elif model_name == 'snail':
    print("HINT: SNAIL works only in PyTorch 0.3.1")
    model = SNAIL(sentence_encoder, N, K)
    framework.train(model, model_name, 25, N, N, K, 1, learning_rate=1e-2, weight_decay=0, optimizer=optim.SGD)
elif model_name == 'metanet':
    model = MetaNet(N, K, train_data_loader.word_vec_mat, max_length)
    framework.train(model, model_name, 1, N, N, K, 1, learning_rate=5e-3, weight_decay=0, optimizer=optim.Adam, train_iter=300000)
else:
    raise NotImplementedError

예제 #3
0
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))
예제 #4
0
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)))