rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Use model.transformer.embed_dim if model.transformer.hidden_size doesn't exist

This commit is contained in:
Gnome Ann 2021-11-10 17:47:14 -05:00
parent 752e19a2bb
commit 11b0291bc4
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
aiserver.py
View File

@ -591,6 +591,12 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
return stopping_criteria
transformers.generation_utils.GenerationMixin._get_stopping_criteria = new_get_stopping_criteria
def get_hidden_size_from_model(model):
try:
return int(model.transformer.hidden_size)
except:
return int(model.transformer.embed_dim)
# If custom GPT Neo model was chosen
if(vars.model == "NeoCustom"):
model_config = open(vars.custmodpth + "/config.json", "r")
@ -632,20 +638,20 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
if(vars.hascuda):
if(vars.usegpu):
model = AutoModelForCausalLM.from_pretrained(vars.model, device=0)
vars.modeldim = int(model.transformer.hidden_size)
vars.modeldim = get_hidden_size_from_model(model)
model = model.to(0)
generator = model.generate
elif(vars.breakmodel): # Use both RAM and VRAM (breakmodel)
model = AutoModelForCausalLM.from_pretrained(vars.model)
vars.modeldim = int(model.transformer.hidden_size)
vars.modeldim = get_hidden_size_from_model(model)
device_config(model)
else:
model = AutoModelForCausalLM.from_pretrained(vars.model)
vars.modeldim = int(model.transformer.hidden_size)
vars.modeldim = get_hidden_size_from_model(model)
generator = model.generate
else:
model = AutoModelForCausalLM.from_pretrained(vars.model)
vars.modeldim = int(model.transformer.hidden_size)
vars.modeldim = get_hidden_size_from_model(model)
generator = model.generate
# Suppress Author's Note by flagging square brackets (Old implementation)