Add Transformers and lm-eval compatibility

chat_template.jinja

@@ -0,0 +1 @@
+{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}

config.json

@@ -1,39 +1,27 @@
 {
-  "model": {
-    "vocab_size": 4096,
-    "max_seq_len": 128,
-    "d_model": 384,
-    "n_layers": 6,
-    "n_heads": 6,
-    "ffn_hidden": 768,
-    "dropout": 0.1,
-    "pad_id": 0,
-    "eos_id": 2
+  "architectures": [
+    "LoafLMForCausalLM"
+  ],
+  "attention_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_loaflm.LoafLMConfig",
+    "AutoModelForCausalLM": "modeling_loaflm.LoafLMForCausalLM"
   },
-  "train": {
-    "input_jsonl": "loaflm_60k.jsonl",
-    "data_dir": "data",
-    "output_dir": "checkpoints",
-    "skip_prepare": false,
-    "vocab_size": 4096,
-    "max_seq_len": 128,
-    "d_model": 384,
-    "n_layers": 6,
-    "n_heads": 6,
-    "ffn_hidden": 768,
-    "dropout": 0.1,
-    "batch_size": 32,
-    "max_steps": 10000,
-    "learning_rate": 0.0003,
-    "min_lr": 3e-05,
-    "warmup_steps": 200,
-    "weight_decay": 0.1,
-    "eval_ratio": 0.02,
-    "eval_batches": 50,
-    "eval_interval": 200,
-    "save_interval": 500,
-    "report_interval": 50,
-    "seed": 42,
-    "cpu": false
-  }
-}
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 768,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 128,
+  "model_type": "loaflm",
+  "num_attention_heads": 6,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.1",
+  "use_cache": false,
+  "vocab_size": 4096
+}

configuration_loaflm.py

@@ -0,0 +1,116 @@
+"""Hugging Face configuration for LoafLM."""
+from __future__ import annotations
+
+from transformers import PreTrainedConfig
+
+
+class LoafLMConfig(PreTrainedConfig):
+    """Configuration for the tiny decoder-only LoafLM architecture."""
+
+    model_type = "loaflm"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 4096,
+        max_position_embeddings: int = 128,
+        hidden_size: int = 384,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 6,
+        intermediate_size: int = 768,
+        hidden_dropout_prob: float = 0.1,
+        attention_dropout_prob: float = 0.1,
+        layer_norm_eps: float = 1e-5,
+        initializer_range: float = 0.02,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = False,
+        use_cache: bool = False,
+        # Old MLX config aliases. These make older local config values easy to reuse.
+        max_seq_len: int | None = None,
+        d_model: int | None = None,
+        n_layers: int | None = None,
+        n_heads: int | None = None,
+        ffn_hidden: int | None = None,
+        dropout: float | None = None,
+        pad_id: int | None = None,
+        eos_id: int | None = None,
+        **kwargs,
+    ) -> None:
+        if max_seq_len is not None:
+            max_position_embeddings = max_seq_len
+        if d_model is not None:
+            hidden_size = d_model
+        if n_layers is not None:
+            num_hidden_layers = n_layers
+        if n_heads is not None:
+            num_attention_heads = n_heads
+        if ffn_hidden is not None:
+            intermediate_size = ffn_hidden
+        if dropout is not None:
+            hidden_dropout_prob = dropout
+            attention_dropout_prob = dropout
+        if pad_id is not None:
+            pad_token_id = pad_id
+        if eos_id is not None:
+            eos_token_id = eos_id
+
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+        if max_position_embeddings <= 0:
+            raise ValueError("max_position_embeddings must be positive")
+
+        self.vocab_size = int(vocab_size)
+        self.max_position_embeddings = int(max_position_embeddings)
+        self.hidden_size = int(hidden_size)
+        self.num_hidden_layers = int(num_hidden_layers)
+        self.num_attention_heads = int(num_attention_heads)
+        self.intermediate_size = int(intermediate_size)
+        self.hidden_dropout_prob = float(hidden_dropout_prob)
+        self.attention_dropout_prob = float(attention_dropout_prob)
+        self.layer_norm_eps = float(layer_norm_eps)
+        self.initializer_range = float(initializer_range)
+        self.use_cache = bool(use_cache)
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            is_encoder_decoder=False,
+            **kwargs,
+        )
+
+    # Aliases retained for LoafLM's original MLX naming.
+    @property
+    def max_seq_len(self) -> int:
+        return self.max_position_embeddings
+
+    @property
+    def d_model(self) -> int:
+        return self.hidden_size
+
+    @property
+    def n_layers(self) -> int:
+        return self.num_hidden_layers
+
+    @property
+    def n_heads(self) -> int:
+        return self.num_attention_heads
+
+    @property
+    def ffn_hidden(self) -> int:
+        return self.intermediate_size
+
+    @property
+    def dropout(self) -> float:
+        return self.hidden_dropout_prob
+
+    @property
+    def pad_id(self) -> int:
+        return self.pad_token_id
+
+    @property
+    def eos_id(self) -> int:
+        return self.eos_token_id

conversion_report.json

@@ -0,0 +1,11 @@
+{
+  "source_checkpoint": "checkpoints/best_model.safetensors",
+  "output_dir": "/Users/guus/Documents/Projects/LoafLM/LoafLM-10M-HF",
+  "source_tensor_count": 77,
+  "hf_tensor_count": 77,
+  "parameter_count_unique": 10298880,
+  "save_reload_max_abs_logit_difference": 0.0,
+  "conversion_notes": [],
+  "lossless": true,
+  "quantized": false
+}

generation_config.json

@@ -0,0 +1,11 @@
+{
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 2,
+  "max_new_tokens": 64,
+  "pad_token_id": 0,
+  "temperature": 0.3,
+  "top_k": 50,
+  "transformers_version": "5.14.1",
+  "use_cache": false
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:76630cfd68fbac0da939063b8945425655f8e8464f00a140f0a654de8755c2d7
-size 41202478
+oid sha256:b314ccfff53e39bbaf2cc4109c69e6e944720a05fcd5c7e112f4a5665d435af9
+size 41202480

modeling_loaflm.py

@@ -0,0 +1,281 @@
+"""Hugging Face Transformers implementation of LoafLM.
+
+This intentionally mirrors the original MLX architecture:
+- learned token and position embeddings
+- pre-norm causal self-attention
+- fused QKV projection
+- ReLU feed-forward network
+- language-model head (tied only when requested by the config)
+"""
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_loaflm import LoafLMConfig
+except ImportError:  # Allows local direct imports during conversion/testing.
+    from configuration_loaflm import LoafLMConfig
+
+
+class LoafAttention(nn.Module):
+    def __init__(self, config: LoafLMConfig) -> None:
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size)
+        self.out = nn.Linear(config.hidden_size, config.hidden_size)
+        self.drop = nn.Dropout(config.attention_dropout_prob)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, seq_len, channels = hidden_states.shape
+
+        qkv = self.qkv(hidden_states).reshape(
+            batch_size,
+            seq_len,
+            3,
+            self.n_heads,
+            self.head_dim,
+        )
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        query, key, value = qkv[0], qkv[1], qkv[2]
+
+        scores = torch.matmul(query, key.transpose(-2, -1)) * self.scale
+
+        causal_mask = torch.ones(
+            (seq_len, seq_len),
+            dtype=torch.bool,
+            device=hidden_states.device,
+        ).tril()
+        combined_mask = causal_mask.view(1, 1, seq_len, seq_len)
+
+        query_mask = None
+        if attention_mask is not None:
+            if attention_mask.ndim != 2:
+                raise ValueError("attention_mask must have shape [batch, sequence]")
+            key_mask = attention_mask[:, None, None, :].to(torch.bool)
+            combined_mask = combined_mask & key_mask
+            query_mask = attention_mask[:, None, :, None].to(scores.dtype)
+
+        # Use a finite minimum. Fully padded query rows then remain numerically safe;
+        # they are zeroed by query_mask below and never influence real tokens.
+        scores = scores.masked_fill(~combined_mask, torch.finfo(scores.dtype).min)
+        attention_probs = torch.softmax(scores, dim=-1)
+        if query_mask is not None:
+            attention_probs = attention_probs * query_mask
+        attention_probs = self.drop(attention_probs)
+
+        context = torch.matmul(attention_probs, value)
+        context = context.permute(0, 2, 1, 3).reshape(batch_size, seq_len, channels)
+        output = self.out(context)
+        return output, attention_probs if output_attentions else None
+
+
+class LoafFFN(nn.Module):
+    def __init__(self, config: LoafLMConfig) -> None:
+        super().__init__()
+        self.up = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.down = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.drop = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.drop(self.down(F.relu(self.up(hidden_states))))
+
+
+class LoafBlock(nn.Module):
+    def __init__(self, config: LoafLMConfig) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.attn = LoafAttention(config)
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.ffn = LoafFFN(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attn_output, attn_weights = self.attn(
+            self.norm1(hidden_states),
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = hidden_states + attn_output
+        hidden_states = hidden_states + self.ffn(self.norm2(hidden_states))
+        return hidden_states, attn_weights
+
+
+class LoafLMPreTrainedModel(PreTrainedModel):
+    config_class = LoafLMConfig
+    base_model_prefix = ""
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["LoafBlock"]
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class LoafLMForCausalLM(LoafLMPreTrainedModel, GenerationMixin):
+    """Decoder-only LoafLM with a causal language-modeling head."""
+
+    # Transformers 5 expects a target -> source mapping here (rather than the
+    # list of target keys accepted by Transformers 4).
+    _tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}
+
+    def __init__(self, config: LoafLMConfig) -> None:
+        super().__init__(config)
+        self.tok_emb = nn.Embedding(
+            config.vocab_size,
+            config.hidden_size,
+            padding_idx=None,  # Matches the original MLX embedding behavior.
+        )
+        self.pos_emb = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.drop = nn.Dropout(config.hidden_dropout_prob)
+        self.blocks = nn.ModuleList([LoafBlock(config) for _ in range(config.num_hidden_layers)])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.tok_emb
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.tok_emb = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position=None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast | tuple:
+        del past_key_values, cache_position, kwargs
+        output_attentions = bool(output_attentions)
+        output_hidden_states = bool(output_hidden_states)
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Provide exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.tok_emb(input_ids)
+        batch_size, seq_len, _ = inputs_embeds.shape
+
+        if seq_len > self.config.max_position_embeddings:
+            raise ValueError(
+                f"Sequence length {seq_len} exceeds LoafLM's context length "
+                f"{self.config.max_position_embeddings}. Truncate the input first."
+            )
+
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_len),
+                dtype=torch.long,
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            # Keeps content positions starting at zero even when lm-eval left-pads a batch.
+            position_ids = attention_mask.long().cumsum(dim=-1) - 1
+            position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+        position_ids = position_ids.clamp(min=0, max=self.config.max_position_embeddings - 1)
+
+        hidden_states = inputs_embeds + self.pos_emb(position_ids)
+        hidden_states = self.drop(hidden_states)
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for block in self.blocks:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            hidden_states, attn_weights = block(
+                hidden_states,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+            )
+            if output_attentions:
+                all_attentions = all_attentions + (attn_weights,)
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        if not return_dict:
+            outputs = (logits, None, all_hidden_states, all_attentions)
+            return ((loss,) + outputs) if loss is not None else outputs
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        # LoafLM is tiny, so generation intentionally recomputes the full 128-token
+        # context instead of maintaining a KV cache.
+        if input_ids.shape[1] > self.config.max_position_embeddings:
+            input_ids = input_ids[:, -self.config.max_position_embeddings :]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -self.config.max_position_embeddings :]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "use_cache": False,
+        }

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|im_start|>",
+  "eos_token": "<|im_end|>",
+  "model_max_length": 128,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}