wehub-resource-sync
108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
import argparse
|
|
|
|
import torch
|
|
|
|
|
|
def fill_config(args):
|
|
# dirty work
|
|
args.device = torch.device(args.gpu)
|
|
args.dec_ninp = args.nhid * 3 if args.title else args.nhid * 2
|
|
args.fnames = [args.train_file, args.valid_file, args.test_file]
|
|
return args
|
|
|
|
|
|
def vocab_config(
|
|
args, ent_vocab, rel_vocab, text_vocab, ent_text_vocab, title_vocab
|
|
):
|
|
# dirty work
|
|
args.ent_vocab = ent_vocab
|
|
args.rel_vocab = rel_vocab
|
|
args.text_vocab = text_vocab
|
|
args.ent_text_vocab = ent_text_vocab
|
|
args.title_vocab = title_vocab
|
|
return args
|
|
|
|
|
|
def get_args():
|
|
args = argparse.ArgumentParser(description="Graph Writer in DGL")
|
|
args.add_argument("--nhid", default=500, type=int, help="hidden size")
|
|
args.add_argument("--nhead", default=4, type=int, help="number of heads")
|
|
args.add_argument("--head_dim", default=125, type=int, help="head dim")
|
|
args.add_argument(
|
|
"--weight_decay", default=0.0, type=float, help="weight decay"
|
|
)
|
|
args.add_argument(
|
|
"--prop", default=6, type=int, help="number of layers of gnn"
|
|
)
|
|
args.add_argument("--title", action="store_true", help="use title input")
|
|
args.add_argument("--test", action="store_true", help="inference mode")
|
|
args.add_argument("--batch_size", default=32, type=int, help="batch_size")
|
|
args.add_argument(
|
|
"--beam_size", default=4, type=int, help="beam size, 1 for greedy"
|
|
)
|
|
args.add_argument("--epoch", default=20, type=int, help="training epoch")
|
|
args.add_argument(
|
|
"--beam_max_len",
|
|
default=200,
|
|
type=int,
|
|
help="max length of the generated text",
|
|
)
|
|
args.add_argument(
|
|
"--enc_lstm_layers",
|
|
default=2,
|
|
type=int,
|
|
help="number of layers of lstm",
|
|
)
|
|
args.add_argument("--lr", default=1e-1, type=float, help="learning rate")
|
|
# args.add_argument('--lr_decay', default=1e-8, type=float, help='')
|
|
args.add_argument("--clip", default=1, type=float, help="gradient clip")
|
|
args.add_argument(
|
|
"--emb_drop", default=0.0, type=float, help="embedding dropout"
|
|
)
|
|
args.add_argument(
|
|
"--attn_drop", default=0.1, type=float, help="attention dropout"
|
|
)
|
|
args.add_argument("--drop", default=0.1, type=float, help="dropout")
|
|
args.add_argument("--lp", default=1.0, type=float, help="length penalty")
|
|
args.add_argument(
|
|
"--graph_enc",
|
|
default="gtrans",
|
|
type=str,
|
|
help="gnn mode, we only support the graph transformer now",
|
|
)
|
|
args.add_argument(
|
|
"--train_file",
|
|
default="data/unprocessed.train.json",
|
|
type=str,
|
|
help="training file",
|
|
)
|
|
args.add_argument(
|
|
"--valid_file",
|
|
default="data/unprocessed.val.json",
|
|
type=str,
|
|
help="validation file",
|
|
)
|
|
args.add_argument(
|
|
"--test_file",
|
|
default="data/unprocessed.test.json",
|
|
type=str,
|
|
help="test file",
|
|
)
|
|
args.add_argument(
|
|
"--save_dataset",
|
|
default="data.pickle",
|
|
type=str,
|
|
help="save path of dataset",
|
|
)
|
|
args.add_argument(
|
|
"--save_model",
|
|
default="saved_model.pt",
|
|
type=str,
|
|
help="save path of model",
|
|
)
|
|
|
|
args.add_argument("--gpu", default=0, type=int, help="gpu mode")
|
|
args = args.parse_args()
|
|
args = fill_config(args)
|
|
return args
|