def init_context_word_embs_model(model_path, device, force_reload=False, temperature=1.0, top_k=None, top_p=None,
optimize=None, silence=True):
global CONTEXT_WORD_EMBS_MODELS
model_name = os.path.basename(model_path)
if model_name in CONTEXT_WORD_EMBS_MODELS and not force_reload:
CONTEXT_WORD_EMBS_MODELS[model_name].device = device
if temperature != 1.0:
CONTEXT_WORD_EMBS_MODELS[model_name].temperature = temperature
if top_k:
CONTEXT_WORD_EMBS_MODELS[model_name].top_k = top_k
if top_p:
CONTEXT_WORD_EMBS_MODELS[model_name].top_p = top_p
CONTEXT_WORD_EMBS_MODELS[model_name].silence = silence
return CONTEXT_WORD_EMBS_MODELS[model_name]
if 'distilbert' in model_path.lower():
model = nml.DistilBert(model_path, device=device, temperature=temperature, top_k=top_k, top_p=top_p, silence=silence)
elif 'roberta' in model_path.lower():
model = nml.Roberta(model_path, device=device, temperature=temperature, top_k=top_k, top_p=top_p, silence=silence)
elif 'bert' in model_path.lower():
model = nml.Bert(model_path, device=device, temperature=temperature, top_k=top_k, top_p=top_p, silence=silence)
elif 'xlnet' in model_path.lower():
model = nml.XlNet(model_path, device=device, temperature=temperature, top_k=top_k, top_p=top_p, optimize=optimize,
silence=silence)
else:
raise ValueError('Model name value is unexpected. Only support BERT, DistilBERT, RoBERTa and XLNet model.')
CONTEXT_WORD_EMBS_MODELS[model_name] = model
return model
def init_context_word_embs_model(model_path,
model_type,
device,
force_reload=False,
batch_size=32,
top_k=None,
silence=True,
use_custom_api=False):
global CONTEXT_WORD_EMBS_MODELS
model_name = '_'.join(
[os.path.basename(model_path), model_type,
str(device)])
if model_name in CONTEXT_WORD_EMBS_MODELS and not force_reload:
CONTEXT_WORD_EMBS_MODELS[model_name].top_k = top_k
CONTEXT_WORD_EMBS_MODELS[model_name].batch_size = batch_size
CONTEXT_WORD_EMBS_MODELS[model_name].silence = silence
return CONTEXT_WORD_EMBS_MODELS[model_name]
if use_custom_api:
if model_type == 'distilbert':
model = nml.DistilBert(model_path,
device=device,
top_k=top_k,
silence=silence,
batch_size=batch_size)
elif model_type == 'roberta':
model = nml.Roberta(model_path,
device=device,
top_k=top_k,
silence=silence,
batch_size=batch_size)
elif model_type == 'bert':
model = nml.Bert(model_path,
device=device,
top_k=top_k,
silence=silence,
batch_size=batch_size)
else:
raise ValueError(
'Model type value is unexpected. Only support bert and roberta models.'
)
else:
if model_type in ['distilbert', 'bert', 'roberta', 'bart']:
model = nml.FmTransformers(model_path,
model_type=model_type,
device=device,
batch_size=batch_size,
top_k=top_k,
silence=silence)
else:
raise ValueError(
'Model type value is unexpected. Only support bert and roberta models.'
)
CONTEXT_WORD_EMBS_MODELS[model_name] = model
return model