mirror of
https://github.com/wassname/Castor.git
synced 2026-09-25 13:10:11 +08:00
212 lines
9.4 KiB
Python
212 lines
9.4 KiB
Python
import argparse
|
|
|
|
import os
|
|
import shlex
|
|
import subprocess
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
|
|
import utils
|
|
from external_features import stopped, stemmed, compute_idf_weighted_overlap, compute_overlap,\
|
|
get_qadata_only_idf, set_external_features_as_per_paper,\
|
|
set_external_features_as_per_paper_and_stem
|
|
from train import Trainer
|
|
from model import QAModel
|
|
|
|
# logging setup
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
ch = logging.StreamHandler()
|
|
ch.setLevel(logging.DEBUG)
|
|
formatter = logging.Formatter('%(levelname)s - %(message)s')
|
|
ch.setFormatter(formatter)
|
|
logger.addHandler(ch)
|
|
|
|
def logargs(func):
|
|
def inner(*args, **kwargs):
|
|
logger.info('%s : %s %s' % (func.__name__, args, kwargs))
|
|
return func(*args, **kwargs)
|
|
return inner
|
|
|
|
|
|
def compute_map_mrr(dataset_folder, set_folder, test_scores, run_name_prefix=None):
|
|
# logger.info("Running trec_eval script...")
|
|
N = len(test_scores)
|
|
|
|
qids_test, y_test = utils.get_test_qids_labels(dataset_folder, set_folder)
|
|
|
|
# Call TrecEval code to calc MAP and MRR
|
|
df_submission = pd.DataFrame(index=np.arange(N), \
|
|
columns=['qid', 'iter', 'docno', 'rank', 'sim', 'run_id'])
|
|
df_submission['qid'] = qids_test
|
|
df_submission['iter'] = 0
|
|
df_submission['docno'] = np.arange(N)
|
|
df_submission['rank'] = 0
|
|
df_submission['sim'] = test_scores
|
|
df_submission['run_id'] = 'smmodel'
|
|
df_submission.to_csv(os.path.join(dataset_folder, 'submission.txt'), \
|
|
header=False, index=False, sep=' ')
|
|
if run_name_prefix:
|
|
df_submission.to_csv('{}.{}.smrun'.format(run_name_prefix, set_folder),\
|
|
header=False, index=False, sep=" ")
|
|
|
|
df_gold = pd.DataFrame(index=np.arange(N), columns=['qid', 'iter', 'docno', 'rel'])
|
|
df_gold['qid'] = qids_test
|
|
df_gold['iter'] = 0
|
|
df_gold['docno'] = np.arange(N)
|
|
df_gold['rel'] = y_test
|
|
df_gold.to_csv(os.path.join(args.dataset_folder, 'gold.txt'), header=False, index=False, sep=' ')
|
|
|
|
# subprocess.call("/bin/sh run_eval.sh '{}'".format(args.dataset_folder), shell=True)
|
|
pargs = shlex.split("/bin/sh run_eval.sh '{}'".format(dataset_folder))
|
|
p = subprocess.Popen(pargs, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
pout, perr = p.communicate()
|
|
|
|
lines = pout.split(b'\n')
|
|
map = float(lines[0].strip().split()[-1])
|
|
mrr = float(lines[1].strip().split()[-1])
|
|
return map, mrr
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ap = argparse.ArgumentParser(description='pytorch port of the SM model', \
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
ap.add_argument('model_outfile', help='file to save final model')
|
|
ap.add_argument('--word_vectors_file', \
|
|
help='NOTE: a cache will be created for faster loading for word vectors',\
|
|
default="../../data/word2vec/aquaint+wiki.txt.gz.ndim=50.bin")
|
|
ap.add_argument('--dataset_folder', help='directory containing train, dev, test sets', \
|
|
default="../../data/TrecQA")
|
|
ap.add_argument('--classes', type=int, default=2)
|
|
# external features related arguments
|
|
ap.add_argument('--no-ext-feats', action="store_true", \
|
|
help="will not include external features in the model")
|
|
ap.add_argument('--paper-ext-feats', action="store_true", default=True, \
|
|
help="external features as per the paper")
|
|
ap.add_argument('--paper-ext-feats-stem', action="store_true", \
|
|
help="external features as per the paper")
|
|
# system arguments
|
|
ap.add_argument('--cuda', action='store_true', help='use CUDA if available')
|
|
ap.add_argument('--num_threads', help="the number of simultaneous processes to run", \
|
|
type=int, default=4)
|
|
# training arguments
|
|
ap.add_argument('--batch_size', type=int, default=1, help="training mini-batch size")
|
|
ap.add_argument('--filter_width', type=int, default=5, help="number of convolution channels")
|
|
ap.add_argument('--eta', help='Initial learning rate', default=0.001, type=float)
|
|
ap.add_argument('--mom', help='SGD Momentum', default=0.0, type=float)
|
|
ap.add_argument('--train', help='switches to train set', action="store_true")
|
|
# epoch related arguments
|
|
ap.add_argument('--epochs', type=int, default=25, help="number of training epochs")
|
|
ap.add_argument('--patience', type=int, default=5, \
|
|
help="if there is no appreciable change in model after <patience> epochs, then stop")
|
|
# debugging arguments
|
|
ap.add_argument('--debug_single_batch', action="store_true", \
|
|
help="will stop program after training 1 input batch")
|
|
ap.add_argument('--num_conv_filters', default=100, type=int, \
|
|
help="the number of convolution channels (lesser is faster)")
|
|
ap.add_argument('--no_loss_reg', help="no loss regularization", action="store_true")
|
|
ap.add_argument('--test_on_each_epoch', action="store_true", \
|
|
help='runs test on each epoch to track final performance')
|
|
ap.add_argument("--skip-training", help="will load pre-trained model", action="store_true")
|
|
ap.add_argument("--run-name-prefix", help="will output train|dev|test runs with provided prefix")
|
|
ap.add_argument("--stop-punct", help='removes punctuation', action="store_true")
|
|
ap.add_argument("--dash-split", help="split words containing hyphens", action="store_true")
|
|
ap.add_argument("--index-for-corpusIDF", help="fetches idf from Index. provide index path. will\
|
|
generate a vocabFile")
|
|
ap.add_argument('--seed', help='Random seed', type=int, default=1234)
|
|
ap.add_argument('--nocudnn', help='Disable the CuDNN backend', action="store_true")
|
|
|
|
args = ap.parse_args()
|
|
|
|
torch.manual_seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
if args.cuda and torch.cuda.is_available():
|
|
torch.cuda.manual_seed(args.seed)
|
|
if args.nocudnn:
|
|
torch.backends.cudnn.enabled = False
|
|
|
|
torch.set_num_threads(args.num_threads)
|
|
|
|
train_set, dev_set, test_set = 'train-all', 'raw-dev', 'raw-test'
|
|
if args.train:
|
|
train_set, dev_set, test_set = 'train', 'clean-dev', 'clean-test'
|
|
|
|
# cache word embeddings
|
|
cache_file = os.path.splitext(args.word_vectors_file)[0] + '.cache'
|
|
utils.cache_word_embeddings(args.word_vectors_file, cache_file)
|
|
|
|
vocab_size, vec_dim = utils.load_embedding_dimensions(cache_file)
|
|
|
|
# instantiate model
|
|
net = QAModel(vec_dim, args.filter_width, args.num_conv_filters, args.no_ext_feats, cuda=args.cuda)
|
|
|
|
# initialize the trainer
|
|
trainer = Trainer(net, args.eta, args.mom, args.no_loss_reg, vec_dim, args.cuda)
|
|
logger.info("Loading input data...")
|
|
# load input data
|
|
trainer.load_input_data(args.dataset_folder, cache_file, train_set, dev_set, test_set)
|
|
logger.info("Setting up external features...")
|
|
# setup external features
|
|
# TODO: remember to update args.* in testing loop below
|
|
if args.paper_ext_feats:
|
|
logger.info("--paper-ext-feats")
|
|
ext_feats_for_splits = \
|
|
set_external_features_as_per_paper(trainer, args.index_for_corpusIDF)
|
|
# ^^ we are saving the features to be used while testing at the end of training
|
|
elif args.paper_ext_feats_stem:
|
|
logger.info("--paper-ext-feats-stem")
|
|
ext_feats_for_splits = \
|
|
set_external_features_as_per_paper_and_stem(trainer, args.index_for_corpusIDF)
|
|
|
|
|
|
if not args.skip_training:
|
|
best_map = 0.0
|
|
best_model = 0
|
|
|
|
for i in range(args.epochs):
|
|
logger.info('------------- Training epoch {} --------------'.format(i+1))
|
|
train_accuracy = trainer.train(train_set, args.batch_size, args.debug_single_batch)
|
|
if args.debug_single_batch: sys.exit(0)
|
|
|
|
dev_scores = trainer.test(dev_set, args.batch_size)
|
|
|
|
dev_map, dev_mrr = compute_map_mrr(args.dataset_folder, dev_set, dev_scores)
|
|
logger.info("------- MAP {}, MRR {}".format(dev_map, dev_mrr))
|
|
|
|
if dev_map - best_map > 1e-3: # new map is better than best map
|
|
best_model = i
|
|
best_map = dev_map
|
|
|
|
QAModel.save(net, args.model_outfile)
|
|
logger.info('Achieved better dev_map ... saved model')
|
|
|
|
if args.test_on_each_epoch:
|
|
test_scores = trainer.test(test_set, args.batch_size)
|
|
map, mrr = compute_map_mrr(args.dataset_folder, test_set, test_scores)
|
|
logger.info("------- MAP {}, MRR {}".format(map, mrr))
|
|
|
|
if (i - best_model) >= args.patience:
|
|
logger.warning('No improvement since the last {} epochs. Stopping training'\
|
|
.format(i - best_model))
|
|
break
|
|
|
|
logger.info(' ------------ Training epochs completed! ------------')
|
|
logger.info('Best dev MAP in training phase = {:.4f}'.format(best_map))
|
|
|
|
trained_model = QAModel.load(args.model_outfile)
|
|
evaluator = Trainer(trained_model, args.eta, args.mom, args.no_loss_reg, vec_dim, args.cuda)
|
|
|
|
for split in [test_set, dev_set]:
|
|
evaluator.load_input_data(args.dataset_folder, cache_file, None, None, split)
|
|
if args.paper_ext_feats or args.paper_ext_feats_stem:
|
|
evaluator.data_splits[split][-1] = ext_feats_for_splits[split]
|
|
#set_external_features_as_per_paper(evaluator)
|
|
split_scores = evaluator.test(split, args.batch_size)
|
|
map, mrr = compute_map_mrr(args.dataset_folder, split, split_scores, args.run_name_prefix)
|
|
logger.info("-------{} MAP {}, MRR {}".format(split, map, mrr))
|