Update modeling_fake_custom.py

config.json

@@ -9,7 +9,7 @@
   "auto_map": {
     "AutoConfig": "configuration_fake_custom.FakeCustomConfig",
     "AutoModel": "modeling_fake_custom.FakeCustomModel",
-    "AutoModelForCausalLM": "modeling_fake_custom.FakeCustomModelForCausalLM"
+    "AutoModelForCausalLM": "modeling_fake_custom.FakeCustomForCausalLM"
   },
   "bos_token_id": 98,
   "embd_pdrop": 0.1,

configuration_fake_custom.py

@@ -1,5 +1,5 @@
-from transformers import GPT2Config
+from transformers import PretrainedConfig
 
 
-class FakeCustomConfig(GPT2Config):
-    model_type = "fakecustom"
+class FakeCustomConfig(PretrainedConfig):
+    model_type = "fakecustom"

modeling_fake_custom.py

@@ -1,9 +1,13 @@
-from transformers import GPT2Model, GPT2LMHeadModel
+from transformers import PreTrainedModel
 from .configuration_fake_custom import FakeCustomConfig
 
-class FakeCustomModel(GPT2Model):
+class FakeCustomPreTrainedModel(PreTrainedModel):
     config_class = FakeCustomConfig
 
 
-class FakeCustomModelForCausalLM(GPT2LMHeadModel):
-    config_class = FakeCustomConfig
+class FakeCustomModel(FakeCustomPreTrainedModel):
+    _auto_class = "AutoModel"
+
+
+class FakeCustomForCausalLM(FakeCustomPreTrainedModel):
+    _auto_class = "AutoModelForCausalLM"