From edd50fc8098f558d2a36d2d95b8833d253595756 Mon Sep 17 00:00:00 2001 From: ebolam Date: Tue, 27 Sep 2022 17:58:51 -0400 Subject: [PATCH] Fix for GPT2 breakmodel in the UI --- aiserver.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/aiserver.py b/aiserver.py index 86e3429b..1ff967b9 100644 --- a/aiserver.py +++ b/aiserver.py @@ -1565,11 +1565,14 @@ 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") - return utils.num_layers(model_config) + + 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 if vars.model_selected == 'OAI':