Use AutoModelForCausalLM for custom models with a model_type

This commit is contained in:
Gnome Ann 2021-10-05 09:45:12 -04:00
parent 18fa742b02
commit 231621e7c2
1 changed files with 1 additions and 1 deletions

View File

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