Fix for GPT2 breakmodel in the UI

This commit is contained in:
ebolam 2022-09-27 17:58:51 -04:00
parent f1d63f61f3
commit edd50fc809

View File

@ -1565,10 +1565,13 @@ def get_layer_count(model, directory=""):
model_config = AutoConfig.from_pretrained(directory, revision=vars.revision, cache_dir="cache")
else:
model_config = AutoConfig.from_pretrained(model, revision=vars.revision, cache_dir="cache")
if ((utils.HAS_ACCELERATE and model_config.model_type != 'gpt2') or model_config.model_type in ("gpt_neo", "gptj", "xglm", "opt")) and not vars.nobreakmodel:
return utils.num_layers(model_config)
else:
return None
else:
return None
def get_oai_models(key):
vars.oaiapikey = key