Example #1
0
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
Example #2
0
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