Use model.transformer.embed_dim if model.transformer.hidden_size doesn't exist
This commit is contained in:
parent
752e19a2bb
commit
11b0291bc4
14
aiserver.py
14
aiserver.py
|
@ -591,6 +591,12 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
|
||||||
return stopping_criteria
|
return stopping_criteria
|
||||||
transformers.generation_utils.GenerationMixin._get_stopping_criteria = new_get_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 custom GPT Neo model was chosen
|
||||||
if(vars.model == "NeoCustom"):
|
if(vars.model == "NeoCustom"):
|
||||||
model_config = open(vars.custmodpth + "/config.json", "r")
|
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.hascuda):
|
||||||
if(vars.usegpu):
|
if(vars.usegpu):
|
||||||
model = AutoModelForCausalLM.from_pretrained(vars.model, device=0)
|
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)
|
model = model.to(0)
|
||||||
generator = model.generate
|
generator = model.generate
|
||||||
elif(vars.breakmodel): # Use both RAM and VRAM (breakmodel)
|
elif(vars.breakmodel): # Use both RAM and VRAM (breakmodel)
|
||||||
model = AutoModelForCausalLM.from_pretrained(vars.model)
|
model = AutoModelForCausalLM.from_pretrained(vars.model)
|
||||||
vars.modeldim = int(model.transformer.hidden_size)
|
vars.modeldim = get_hidden_size_from_model(model)
|
||||||
device_config(model)
|
device_config(model)
|
||||||
else:
|
else:
|
||||||
model = AutoModelForCausalLM.from_pretrained(vars.model)
|
model = AutoModelForCausalLM.from_pretrained(vars.model)
|
||||||
vars.modeldim = int(model.transformer.hidden_size)
|
vars.modeldim = get_hidden_size_from_model(model)
|
||||||
generator = model.generate
|
generator = model.generate
|
||||||
else:
|
else:
|
||||||
model = AutoModelForCausalLM.from_pretrained(vars.model)
|
model = AutoModelForCausalLM.from_pretrained(vars.model)
|
||||||
vars.modeldim = int(model.transformer.hidden_size)
|
vars.modeldim = get_hidden_size_from_model(model)
|
||||||
generator = model.generate
|
generator = model.generate
|
||||||
|
|
||||||
# Suppress Author's Note by flagging square brackets (Old implementation)
|
# Suppress Author's Note by flagging square brackets (Old implementation)
|
||||||
|
|
Loading…
Reference in New Issue