Upload README.md with huggingface_hub

README.md

@@ -1,37 +1,63 @@
----
-{}
----
-# One-Layer Simple Transformer
-
-A 1-layer transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html).
+1-layer simple transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html).
+Load with
+```python
+class OneLayerTransformer(PreTrainedModel):
+    config_class = LlamaConfig
+    
+    def __init__(self, config: LlamaConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+
+        # Single self-attention layer
+        self.self_attn = nn.MultiheadAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=getattr(config, 'attention_dropout', 0.0),
+            batch_first=True,
+        )
+
+        # Output head
+        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):
+        batch_size, seq_len = input_ids.shape
+
+        # Embeddings
+        hidden_states = self.embed_tokens(input_ids)
+        assert hidden_states.shape == (batch_size, seq_len, self.config.hidden_size)
+
+        # Create causal mask for self-attention
+        causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
+        causal_mask = causal_mask.to(hidden_states.device)
+
+        # Self-attention with residual connection
+        attn_output, _ = self.self_attn(
+            hidden_states,
+            hidden_states,
+            hidden_states,
+            attn_mask=causal_mask,
+            key_padding_mask=None if attention_mask is None else ~attention_mask.bool(),
+        )
+        hidden_states = hidden_states + attn_output
+        assert hidden_states.shape == (batch_size, seq_len, self.config.hidden_size)
+
+        # Output projection
+        logits = self.lm_head(hidden_states)
+        assert logits.shape == (batch_size, seq_len, self.config.vocab_size)
+
+        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}
 
-## Usage
 
-```python
-from transformers import LlamaConfig
-from migrate_models import OneLayerTransformer
 
-# Load the model
 model = OneLayerTransformer.from_pretrained('Butanium/simple-stories-one-layer-simple-transformer')
-
-# Or create from config
-config = LlamaConfig(vocab_size=4096, hidden_size=128, num_hidden_layers=1)
-model = OneLayerTransformer(config)
 ```
-
-## Model Architecture
-
-This model consists of:
-- Token embeddings
-- Single self-attention layer with residual connection
-- Linear output head
-
-It serves as a minimal transformer for understanding attention mechanisms and transformer circuits.
-
-## Training Details
-
-- Trained on SimpleStories dataset
-- Vocabulary size: 4096
-- Hidden size: 128
-- Single self-attention layer
-- 4 attention heads
+The model is trained on the SimpleStories dataset.