Upload model

config.json

@@ -0,0 +1,18 @@
+{
+  "architectures": [
+    "SimpleStories4MModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "config_4m.SimpleStories4MConfig",
+    "AutoModel": "model_4m.SimpleStories4MModel"
+  },
+  "block_size": 1080,
+  "dropout": 0.1,
+  "model_type": "simple_stories_4m",
+  "n_embed": 256,
+  "n_heads": 2,
+  "n_layers": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.2",
+  "vocab_size": 2048
+}

config_4m.py

@@ -0,0 +1,22 @@
+from transformers import PretrainedConfig
+
+class SimpleStories4MConfig(PretrainedConfig):
+    model_type = "simple_stories_4m"
+    
+    def __init__(
+        self, 
+        vocab_size: int = 2048,
+        block_size: int = 1080,
+        n_embed: int = 256,
+        n_heads: int = 2,
+        n_layers: int = 4,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.n_embed = n_embed
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.dropout = dropout
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e8c4a3f27258d9726c7d463899171c9ffe0b8f69de03f7fc3edf4d39806f403
+size 55267264

model_4m.py

@@ -0,0 +1,22 @@
+from transformers import PreTrainedModel
+from simple_stories_4m_model.config_4m import SimpleStories4MConfig
+from simple_stories_4m_model.nano_gpt_model import NanoGPT
+
+class SimpleStories4MModel(PreTrainedModel):
+    config_class = SimpleStories4MConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        hyperparameters = {
+            "vocab_size": config.vocab_size,
+            "block_size": config.block_size,
+            "n_embed": config.n_embed,
+            "n_heads": config.n_heads,
+            "n_layers": config.n_layers,
+            "dropout": config.dropout,
+        
+        }
+        self.model = NanoGPT(hyperparameters)
+        
+    def forward(self, tensor, targets=None):
+        return self.model(tensor, targets)