Upload AbcTransformer

config.json

@@ -2,6 +2,9 @@
   "architectures": [
     "AbcTransformer"
   ],
+  "auto_map": {
+    "AutoModelForCausalLM": "transformers_model.AbcTransformer"
+  },
   "block_size": 128,
   "device": "cpu",
   "dropout": 0.2,

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd54b4b9d81b461d9df19ba3e73ec77388febe5eb8d01a20a5870a9e51cfe897
+oid sha256:682aafa0732ff611771441cd3059543c3b9fba5be2c0f6a0f851cc37baa8f075
 size 18965

transformers_model.py

@@ -0,0 +1,48 @@
+import transformers
+
+import model
+
+
+class AbcTransformerConfig(transformers.PretrainedConfig):
+    model_type = 'abc-transformer'
+    def __init__(
+            self,
+            vocab_size=113,
+            n_embd=384,
+            block_size=128,
+            n_heads=6,
+            n_layers=6,
+            dropout=0.2,
+            device=None,
+            **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.block_size = block_size
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.dropout = dropout
+        self.device = device
+        super().__init__(**kwargs)
+
+class AbcTransformer(transformers.PreTrainedModel):
+    config_class = AbcTransformerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = model.AbcTransformer(
+            vocab_size=config.vocab_size,
+            n_embd=config.n_embd,
+            block_size=config.block_size,
+            n_heads=config.n_heads,
+            n_layers=config.n_layers,
+            dropout=config.dropout,
+            device=config.device,
+        )
+    
+    def forward(self, tensor, labels):
+        return self.model(tensor, labels)
+
+transformers.AutoConfig.register('abc-transformer', AbcTransformerConfig)
+AbcTransformer.register_for_auto_class("AutoModelForCausalLM")
+transformers.AutoModel.register(AbcTransformerConfig, AbcTransformer)