Use AutoModelForCausalLM for custom models with a model_type

2021-10-05 09:45:12 -04:00 · 2021-10-05 09:45:12 -04:00 · 231621e7c2
parent 18fa742b02
commit 231621e7c2
1 changed files with 1 additions and 1 deletions
--- a/aiserver.py
+++ b/aiserver.py
@ -487,7 +487,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
            model_config = open(vars.custmodpth + "/config.json", "r")
            js   = json.load(model_config)
            if("model_type" in js):
-                model     = vars.custmodpth
+                model     = AutoModelForCausalLM.from_pretrained(vars.custmodpth)
            else:
                model     = GPTNeoForCausalLM.from_pretrained(vars.custmodpth)
            tokenizer = GPT2Tokenizer.from_pretrained(vars.custmodpth)