def args_parse(config_file=''):
parser = argparse.ArgumentParser(description="fast-bbdl")
parser.add_argument(
"--config_file", default="", help="path to config file", type=str
)
parser.add_argument("--opts", help="Modify config options using the command-line key value", default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
config_file = args.config_file or config_file
if config_file != "":
cfg.merge_from_file(get_abs_path('configs', config_file))
cfg.merge_from_list(args.opts)
cfg.freeze()
name = cfg.MODEL.NAME
output_dir = cfg.OUTPUT_DIR
logger = setup_logger(name, get_abs_path(output_dir), 0)
logger.info(args)
if config_file != '':
logger.info("Loaded configuration file {}".format(config_file))
with open(get_abs_path('configs', config_file), 'r') as cf:
config_str = "\n" + cf.read()
logger.info(config_str)
logger.info("Running with config:\n{}".format(cfg))
return cfg
def load_model(args):
from bbcm.config import cfg
cfg.merge_from_file(get_abs_path('configs', args.config_file))
tokenizer = BertTokenizer.from_pretrained(cfg.MODEL.BERT_CKPT)
file_dir = get_abs_path("checkpoints", cfg.MODEL.NAME)
if cfg.MODEL.NAME in ['bert4csc', 'macbert4csc']:
model = BertForCsc.load_from_checkpoint(os.path.join(
file_dir, args.ckpt_fn),
cfg=cfg,
tokenizer=tokenizer)
else:
model = SoftMaskedBertModel.load_from_checkpoint(os.path.join(
file_dir, args.ckpt_fn),
cfg=cfg,
tokenizer=tokenizer)
model.eval()
model.to(cfg.MODEL.DEVICE)
return model
def load_model_directly(ckpt_file, config_file):
# Example:
# ckpt_fn = 'SoftMaskedBert/epoch=02-val_loss=0.02904.ckpt' (find in checkpoints)
# config_file = 'csc/train_SoftMaskedBert.yml' (find in configs)
from bbcm.config import cfg
cp = get_abs_path('checkpoints', ckpt_file)
cfg.merge_from_file(get_abs_path('configs', config_file))
tokenizer = BertTokenizer.from_pretrained(cfg.MODEL.BERT_CKPT)
if cfg.MODEL.NAME in ['bert4csc', 'macbert4csc']:
model = BertForCsc.load_from_checkpoint(cp,
cfg=cfg,
tokenizer=tokenizer)
else:
model = SoftMaskedBertModel.load_from_checkpoint(cp,
cfg=cfg,
tokenizer=tokenizer)
model.eval()
model.to(cfg.MODEL.DEVICE)
return model