Upload README.md with huggingface_hub

README.md

@@ -1,29 +1,32 @@
----
-{}
----
-0-layer transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html).Load with ```python
-class ZeroLayerTransformer(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, config):
-        super().__init__()
-        self.config = 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)
-
-    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)
-
-        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)
-            )
-
-        return {"loss": loss, "logits": logits}
+# 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
+
+```python
+from transformers import LlamaConfig
+from migrate_models import ZeroLayerTransformer
+
+# Load the model
+model = ZeroLayerTransformer.from_pretrained('Butanium/simple-stories-zero-layer-simple-transformer')
+
+# Or create from config
+config = LlamaConfig(vocab_size=4096, hidden_size=128, num_hidden_layers=0)
+model = ZeroLayerTransformer(config)
 ```
-model = ZeroLayerTransformer.from_pretrained('Butanium/simple-stories-zero-layer-simple-transformer')
+
+## Model Architecture
+
+This model consists of only:
+- Token embeddings
+- Linear output head (no transformer layers)
+
+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)