def __init__(self, args, state_dict=None):
self.args = args
self.updates = 0
# -----------------------------------------------
# select model
network = networks.get_class(args.model_class)
# -----------------------------------------------
# select config
if args.model_class in ['bert2tag', 'roberta2tag']:
args.num_labels = len(Idx2Tag)
else:
args.num_labels = 2
logger.info('Config num_labels = %d' % args.num_labels)
model_config = config_class[args.pretrain_model_type].from_pretrained(
args.cache_dir, num_labels=args.num_labels)
# -----------------------------------------------
# load network
self.network = network.from_pretrained(args.cache_dir,
config=model_config)
if state_dict is not None:
self.network.load_state_dict(state_dict)
logger.info('loaded pretrain model state_dict')
def __init__(self, args, state_dict=None):
self.args = args
self.updates = 0
# select model
network = networks.get_class(args)
# select config
args.num_labels = 2 if args.model_class != 'bert2tag' else len(Idx2Tag)
logger.info('Config num_labels = %d' %args.num_labels)
model_config = config_class[args.pretrain_model_type].from_pretrained(args.cache_dir, num_labels=args.num_labels)
# load pretrained model
self.network = network.from_pretrained(args.cache_dir, config=model_config)
# load checkpoint
if state_dict is not None:
self.network.load_state_dict(state_dict)
logger.info('loaded checkpoint state_dict')