Update aiserver.py

Better way of checking for the model type
This commit is contained in:
henk717 2021-09-25 16:50:24 +02:00
parent 7d35f825c6
commit c9290d02dc
1 changed files with 1 additions and 1 deletions

View File

@ -404,7 +404,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
if(vars.model == "NeoCustom"):
model_config = open(vars.custmodpth + "/config.json", "r")
js = json.load(model_config)
if("architectures" in js):
if("model_type" in js):
model = vars.custmodpth
else:
model = GPTNeoForCausalLM.from_pretrained(vars.custmodpth)