def __init__(self, dataset="django"):
arg_parser = init_arg_parser()
if dataset == "django":
args = init_arg_parser().parse_args("--mode test \
--load_model src/external_repos/tranX/data/pretrained_models/django.bin \
--beam_size 15 \
--test_file src/external_repos/tranX/data/django/test.bin \
--save_decode_to 0.test.decode \
--decode_max_time_step 100 \
--example_preprocessor django_example_processor"
.split())
elif dataset == "conala":
args = init_arg_parser().parse_args("--mode test \
--load_model src/external_repos/tranX/data/pretrained_models/conala.bin \
--beam_size 15 \
--test_file src/external_repos/tranX/data/conala/test.bin \
--save_decode_to 0.test.decode \
--decode_max_time_step 100 \
--example_preprocessor conala_example_processor"
.split())
self.parser = StandaloneParser(args.parser,
args.load_model,
args.example_preprocessor,
beam_size=args.beam_size,
cuda=args.cuda)
def main(args_string=None):
arg_parser = init_arg_parser()
args = init_config(arg_parser, args_string)
if args.mode == 'train':
train(args)
elif args.mode == 'test':
test(args)
else:
raise RuntimeError('unknown mode')
def load(model_path, cuda=False):
decoder_params = torch.load(model_path,
map_location=lambda storage, loc: storage)
decoder_params['args'].cuda = cuda
# update saved args
saved_args = decoder_params['args']
update_args(saved_args, init_arg_parser())
model = ParaphraseIdentificationModel(
saved_args, decoder_params['vocab'],
decoder_params['transition_system'])
model.load_state_dict(decoder_params['state_dict'])
if cuda: model = model.cuda()
model.eval()
return model
def load(cls, model_path, cuda=False):
params = torch.load(model_path, map_location=lambda storage, loc: storage)
vocab = params['vocab']
transition_system = params['transition_system']
saved_args = params['args']
# update saved args
update_args(saved_args, init_arg_parser())
saved_state = params['state_dict']
saved_args.cuda = cuda
parser = cls(saved_args, vocab, transition_system)
parser.load_state_dict(saved_state)
if cuda: parser = parser.cuda()
parser.eval()
return parser
parser_cls = Registrable.by_name(args.parser)
model = parser_cls.load(model_path=args.load_model, cuda=args.cuda)
decode_results = []
count = 0
hyps = model.parse(beam_size)
decoded_hyps = []
for hyp_id, hyp in enumerate(hyps):
try:
hyp.code = model.transition_system.ast_to_surface_code(hyp.tree)
print(hyp.code)
decoded_hyps.append(hyp)
except:
pass
"""
if args.save_decode_to:
pickle.dump(decode_results, open(args.save_decode_to, 'wb'))
"""
if __name__ == '__main__':
arg_parser = init_arg_parser()
args = init_config()
print(args, file=sys.stderr)
if args.mode == 'train':
train(args)
elif args.mode == 'test':
test(args)
else:
raise RuntimeError('unknown mode')
