ForCausalLM

README.md

@@ -1,3 +1,5 @@
 ---
 license: apache-2.0
 ---
+
+This is the same as `test_dynamic_model`, but with a `generate`-compatible class

config.json

@@ -6,7 +6,7 @@
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
     "AutoConfig": "configuration.NewModelConfig",
-    "AutoModel": "modeling.NewModel"
+    "AutoModelForCausalLM": "modeling.NewModelForCausalLM"
   },
   "classifier_dropout": null,
   "hidden_act": "gelu",
@@ -23,7 +23,7 @@
   "pad_token_id": 0,
   "position_embedding_type": "absolute",
   "torch_dtype": "float32",
-  "transformers_version": "4.16.0.dev0",
+  "transformers_version": "4.45.0.dev0",
   "type_vocab_size": 2,
   "use_cache": true,
   "vocab_size": 30522

modeling.py

@@ -1,11 +1,11 @@
 import torch
-from transformers import BertModel
+from transformers import BertLMHeadModel
 
 from .configuration import NewModelConfig
 
-class NewModel(BertModel):
+class NewModelForCausalLM(BertLMHeadModel):
     config_class = NewModelConfig
 
     def __init__(self, config):
         super().__init__(config)
-        self.last_layer = torch.nn.Linear(config.hidden_size, config.new_hidden_size)
+        self.last_layer = torch.nn.Linear(config.hidden_size, config.new_hidden_size)