def main():
parser = argparse.ArgumentParser()
parser.add_argument('--data_path',
type=str,
default=None,
help='Path to train data')
parser.add_argument('--batch_size',
dest='batch_size',
default=64,
type=int)
parser.add_argument('--rbt_model', type=str, default='bert-base-uncased')
parser.add_argument('--got_model', type=str, default='bert-base-uncased')
parser.add_argument('--do_lower_case', action='store_true')
parser.add_argument("--mc_steps",
type=int,
default=100,
help="mc sample steps")
parser.add_argument("--agg_to_file",
type=str,
default=None,
help='path to aggregated file')
parser.add_argument('--tinit',
type=float,
default=3e-2,
help='initial temperature')
parser.add_argument('--C', type=float, default=3e-4, help='scale of temp')
parser.add_argument("--K",
type=int,
default=100,
help="sample subset size")
parser.add_argument("--tk_layer",
type=int,
default=8,
help="nth layer for sent bertscore")
parser.add_argument("--sent_layer",
type=int,
default=1,
help="nth layer for keyword bertscore")
parser.add_argument("--do_copy", action='store_true', help="")
parser.add_argument(
'--alpha',
type=int,
default=3,
help='power for bert score (sentence level semantic sim)')
parser.add_argument(
"--beta",
type=int,
default=8,
help="power for bert_kw in LM for sentence candidate selection")
parser.add_argument('--lm_model_path', type=str, default=None)
parser.add_argument('--lm_vocab_path', type=str, default=None)
parser.add_argument('--max_length',
type=int,
default=40,
help='maybe for lm')
parser.add_argument('--fm_wght',
type=float,
default=0.125,
help='weight of formal')
parser.add_argument('--lm_wght',
type=float,
default=0.125,
help='weight of lm')
parser.add_argument('--temp', type=float, default=8., help='')
parser.add_argument('--rule_path', type=str, default='')
parser.add_argument('--score_mode',
type=str,
default='min',
choices=['sasm', 'sakm', 'smkm'],
help='')
parser.add_argument('--sa_policy',
type=str,
default='one',
choices=['one', 'bs'],
help='')
parser.add_argument('--use_lm_diff', action='store_true')
parser.add_argument('--use_cls_diff', action='store_true')
parser.add_argument('--rand_pos', dest='rand_pos', action='store_true')
parser.add_argument('--no_rand_pos', dest='rand_pos', action='store_false')
parser.set_defaults(rand_pos=True)
parser.add_argument('--nplh', type=int, default=2)
args = parser.parse_args()
gpt_config = GPT2Config.from_pretrianed(args.gpt_model)
gpt_tknzr = GPT2Tokenizer.from_pretrianed(args.gpt_model,
do_lower_case=args.do_lower_case)
gpt_model = GPT2LMHeadModel.from_pretrianed(args.gpt_model,
config=gpt_config,
cache_dir=None)
gpt_model.eval()
rbt_config = RobertaConfig.from_pretrianed(args.rbt_model)
rbt_tknzr = RobertaTokenizer.from_pretrianed(
args.rbt_config, do_lower_case=args.do_lower_case)
rbt_model = RobertaForMTL.from_pretrianed(args.rbt_model,
config=rbt_config)
rbt_model.eval()
scorer = sa_score_funcs.SAScorer(args.batch_size,
rbt_model,
rbt_tknzr,
punc_ids=PUNC_IDS,
gpt_model=gpt_model,
gpt_tknzr=gpt_tknzr,
stopwords=STOPWORDS,
verbose=False,
tinit=args.tinit,
C=args.C,
sent_layer=args.tk_layer,
kw_layer=args.sent_layer,
alpha=args.alpha,
beta=args.beta,
fm_wght=args.fm_wght,
lm_wght=args.lm_wght,
score_mode=args.score_mode,
use_lm_diff=args.use_lm_diff,
use_cls_diff=args.use_cls_diff)
seacher = sa_search.SASearcher(scorer,
rbt_model=rbt_model,
rbt_tokenizer=rbt_tknzr,
k=args.K,
do_copy=args.do_copy,
device=device,
temp=args.temp,
rand_pos=args.rand_pos,
if_fm_voc=None,
nplh=args.nplh)
train_mc(args, seacher, scorer)
