| 0-layer transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html). | |
| Load with | |
| ```python | |
| 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) | |
| 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} | |
| model = ZeroLayerTransformer.from_pretrained('Butanium/simple-stories-zero-layer-simple-transformer') | |
| ``` | |
| The model is trained on the SimpleStories dataset. |