Upload README.md with huggingface_hub

README.md

@@ -1,35 +1,31 @@
----
-{}
----
-# Zero-Layer Simple Transformer
-
-A 0-layer transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html).
-
-## Usage
-
+0-layer transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html).
+Load with
 ```python
-from transformers import LlamaConfig
-from migrate_models import ZeroLayerTransformer
-
-# Load the model
-model = ZeroLayerTransformer.from_pretrained('Butanium/simple-stories-zero-layer-simple-transformer')
+class ZeroLayerTransformer(PreTrainedModel):
+    config_class = LlamaConfig
+    
+    def __init__(self, config: LlamaConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
-# Or create from config
-config = LlamaConfig(vocab_size=4096, hidden_size=128, num_hidden_layers=0)
-model = ZeroLayerTransformer(config)
-```
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        hidden_states = self.embed_tokens(input_ids)
+        logits = self.lm_head(hidden_states)
 
-## Model Architecture
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
 
-This model consists of only:
-- Token embeddings
-- Linear output head (no transformer layers)
+        return {"loss": loss, "logits": logits}
 
-It serves as a baseline for understanding transformer circuits and the importance of attention layers.
 
-## Training Details
 
-- Trained on SimpleStories dataset
-- Vocabulary size: 4096
-- Hidden size: 128
-- No transformer layers (0-layer architecture)
+model = ZeroLayerTransformer.from_pretrained('Butanium/simple-stories-zero-layer-simple-transformer')
+```
+The model is trained on the SimpleStories dataset.