Add Transformers support (modeling + config code), tokenizer/generation configs and model card

README.md

@@ -0,0 +1,196 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+language:
+- en
+tags:
+- comet
+- openmodel-1
+- causal-lm
+- base-model
+- hybrid-architecture
+- custom_code
+datasets:
+- HuggingFaceFW/fineweb-edu
+- HuggingFaceFW/fineweb
+- HuggingFaceTB/finemath
+- Cyrile/dataset-the-stack-v2-dedup-sub
+---
+
+# Comet-1-Preview
+
+Comet-1-Preview is a 7.0B-parameter causal language model from OpenModel-Labs, trained from scratch with the single-file pipeline in [`train.py`](./train.py) (internal name: OpenModel-1). Three out of every four blocks are grouped-query attention; the fourth replaces attention with a gated, dilated causal convolution.
+
+**This checkpoint is a base model.** It continues text; it does not follow chat turns. Prompt it with the beginning of the document you want, or with a few-shot pattern. See [Prompting](#prompting).
+
+It is a preview release: no benchmark results have been published for it, and no license has been declared for the weights yet.
+
+| | |
+|---|---|
+| Parameters | 7,001,436,928 |
+| Weights | `model.safetensors`, float32, 28 GB |
+| Context | trained on 4,096-token sequences; position table covers 8,192 |
+| Tokenizer | GPT-NeoX-20B BPE, 50,277 tokens |
+| Language | English web text, with maths and source code |
+
+## Quick start
+
+The architecture is not built into Transformers, so loading needs `trust_remote_code=True`. That runs [`modeling_comet.py`](./modeling_comet.py) and [`configuration_comet.py`](./configuration_comet.py) from this repository; they are short and worth reading first.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "OpenModel-Labs/Comet-1-Preview"
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo,
+    trust_remote_code=True,
+    dtype=torch.bfloat16,   # on older 4.x releases that reject `dtype`, use torch_dtype=
+    device_map="auto",      # needs `accelerate`
+)
+
+inputs = tokenizer("Photosynthesis is the process by which", return_tensors="pt").to(model.device)
+output = model.generate(
+    **inputs,
+    max_new_tokens=128,
+    do_sample=True,
+    temperature=0.7,
+    top_p=0.9,
+    repetition_penalty=1.1,
+)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+The sampling values above are ordinary starting points, not tuned settings.
+
+Memory: the weights take about 14 GB in bfloat16 and 28 GB in float32. Training computed in bfloat16, so bfloat16 inference matches the precision the weights were trained under. The generation cache adds roughly 0.6 GB per sequence at 4,096 tokens in bfloat16.
+
+Requirements: a recent PyTorch 2.x and `transformers` 4.46 or newer.
+
+## Prompting
+
+The model was trained on plain documents separated by `<|endoftext|>`. Write the start of a document and let it continue.
+
+These are greedy continuations (`do_sample=False`, 20 new tokens) from the released weights. `\n` marks a line break.
+
+| Prompt | Continuation |
+|---|---|
+| `The capital of Australia is` | ` Canberra.\n\nThe population of Australia is about 23 million.\n\nThe currency of` |
+| `Photosynthesis is the process by which` | ` green plants convert carbon dioxide and water into glucose and oxygen. This process is essential for the survival of` |
+| `def is_prime(n):\n    """Return True if n is a prime number."""\n` | `    if n < 2:\n        return False\n    for i in range(2, int` |
+| `Question: What is 12 + 30?\nAnswer:` | ` 12 + 30 = 42\n\nQuestion: What is 12 + 30 + 42?\nAnswer:` |
+
+The last row shows the usual base-model behaviour: after answering, it writes the next question itself. Stop on a string of your choosing (for example `"\nQuestion:"`) or cap `max_new_tokens`.
+
+Things to know:
+
+- The tokenizer adds no special tokens. `<|endoftext|>` (id 0) is the document separator and the end-of-sequence token for generation. `<|padding|>` (id 1) is the pad token, and batches are padded on the left.
+- No chat template is shipped. `train.py` defines later fine-tuning stages that use a `<|user|>` / `<|assistant|>` turn format, but this checkpoint does not respond to it. Given `<|user|>\nWhat is the capital of France?\n<|endoftext|><|assistant|>\n`, its greedy continuation was `\n# How to Calculate the Area of a Circle\n\nThe area of a circle is the amount`, which is the start of an unrelated document.
+- For question answering or instructions, use a few-shot pattern such as the `Question:` / `Answer:` one above.
+
+## Architecture
+
+```
+tokens -> embedding (50,277 x 4,096)
+       -> 36 blocks, repeating [global, global, global, recall]
+       -> RMSNorm -> logits through the same embedding matrix (tied)
+
+block:   x = x + mixer_scale * mixer(RMSNorm(x))
+         x = x + mlp_scale   * SwiGLU(RMSNorm(x))
+```
+
+| Component | Detail |
+|---|---|
+| Blocks | 36 in total: 27 `global`, 9 `recall` (blocks 3, 7, 11, ..., 35) |
+| Width | 4,096 |
+| `global` mixer | Causal grouped-query attention: 32 query heads, 8 key/value heads, head size 128. Queries and keys pass through a per-head RMSNorm, then RoPE (base 1,000,000). |
+| `recall` mixer | Project 4,096 to a gate and a value of 2,048 channels each; run a depthwise causal convolution over the value (kernel 7, dilation 8); multiply by the sigmoid of the gate; project back to 4,096. |
+| Feed-forward | SwiGLU, inner width 12,288, in every block |
+| Residual scaling | A learned per-channel scale on each branch, initialised to 0.5 |
+| Normalisation | Pre-norm RMSNorm, epsilon 1e-6 |
+| Biases | None |
+
+A `recall` block sees seven positions: the current token and the tokens 8, 16, 24, 32, 40 and 48 places back. The config lists `memory_dilations: [1, 2, 4, 8]`, and a block picks its dilation by its index modulo four. Because `recall` blocks always sit at an index that is 3 modulo 4, all nine use dilation 8 in this checkpoint.
+
+A `global` block has about 193M parameters and a `recall` block about 176M; the embedding matrix has 206M.
+
+## Training
+
+This section describes the settings in `train.py`. It is not a log of the run: the upload does not include the run metadata, so the number of tokens this checkpoint actually saw is not recorded here.
+
+**Data.** Pretraining streams FineWeb-Edu and FineWeb, then gradually mixes in maths and code. The mixture moves linearly from the first column to the second over 12B tokens.
+
+| Source | Start | After anneal | License note in `train.py` |
+|---|---|---|---|
+| [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) (`sample-100BT`) | 70% | 55% | ODC-By, Common Crawl-derived |
+| [HuggingFaceFW/fineweb](https://huggingface.co/datasets/HuggingFaceFW/fineweb) (`sample-100BT`) | 30% | 20% | ODC-By, Common Crawl-derived |
+| [HuggingFaceTB/finemath](https://huggingface.co/datasets/HuggingFaceTB/finemath) (`finemath-4plus`) | 0% | 15% | ODC-By, Common Crawl-derived |
+| [Cyrile/dataset-the-stack-v2-dedup-sub](https://huggingface.co/datasets/Cyrile/dataset-the-stack-v2-dedup-sub) (Python 6%, JavaScript 2%, Java 1%, C++ 0.75%, C 0.25%) | 0% | 10% | per-file license review required |
+
+**Optimisation.**
+
+| Setting | Value |
+|---|---|
+| Sequence length | 4,096 tokens; documents are packed end to end with the end-of-text token between them |
+| Hardware | 8 GPUs, DDP |
+| Batch | 1 sequence per GPU, 16 accumulation steps: 524,288 tokens per optimizer step |
+| Optimizer | AdamW, betas (0.9, 0.95), weight decay 0.1 (not applied to norm gains or residual scales) |
+| Learning rate | 3e-4 peak, linear warm-up over 600M tokens, cosine decay to 3e-5 |
+| Gradient clipping | 1.0 |
+| Precision | float32 parameters, bfloat16 autocast |
+| Token budget | 22 tokens per parameter (about 154B tokens), with an 11-day wall-clock cap on the stage |
+
+`train.py` also contains supervised fine-tuning and SimPO preference-tuning stages. As shown under [Prompting](#prompting), this checkpoint behaves as the output of the pretraining stage.
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `model.safetensors` | Weights, exactly as written by `train.py` (flat keys: `embed.weight`, `blocks.N.*`, `norm.weight`) |
+| `config.json` | Architecture hyperparameters and the `auto_map` that points Transformers at the code below |
+| `configuration_comet.py` | `CometConfig` |
+| `modeling_comet.py` | `CometModel` (hidden states) and `CometForCausalLM` (logits, loss, `generate`) |
+| `generation_config.json` | End-of-sequence and padding ids |
+| `tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json` | Tokenizer (same `tokenizer.json` content as `EleutherAI/gpt-neox-20b`) |
+| `train.py` | The training pipeline, including the reference `OpenModel1` module |
+
+## Implementation notes
+
+`modeling_comet.py` keeps the parameter names and shapes of the training module, so the checkpoint loads with no key remapping. It adds what the training loop did not need:
+
+- **Generation cache.** `use_cache=True` returns a standard `DynamicCache`. Attention blocks store keys and values in it. `recall` blocks store their convolution input in the key slot, so beam search, batch expansion and cache cropping work on them unchanged.
+- **Padding.** An `attention_mask` hides padded keys in attention, zeroes padded positions in the convolution, and numbers real tokens from position 0. A left-padded batch therefore gives the same logits as running each row alone.
+- **Loss.** Pass `labels=input_ids`; the shift happens inside the model and `-100` is ignored.
+- **Gradient checkpointing.** `model.gradient_checkpointing_enable()`.
+
+Not supported: `output_attentions` (attention runs through `scaled_dot_product_attention`), legacy tuple caches, and inputs longer than 8,192 positions, which raise an error.
+
+## Verification
+
+What was checked before these files were published:
+
+- **Layout.** A full-size `CometForCausalLM` has the same 353 tensor names and shapes as the header of the released `model.safetensors`.
+- **Equivalence.** With identical weights on small random configurations, logits, hidden states and loss match the `OpenModel1` module from `train.py` exactly in float32 (maximum absolute difference 0).
+- **Cache and padding.** Token-by-token and chunked decoding match a full forward pass to about 1e-6. Greedy `generate` with the cache matches recomputing from scratch, beam search matches with and without the cache, and left-padded batches match unbatched runs.
+- **Released weights.** A float32 CPU forward pass over the full checkpoint gave perplexity 4.5 on a 68-token encyclopedic paragraph and 1.7 on a short Python function, put 98.8% on ` Tokyo` after a list of capitals ending in `The capital of Japan is`, and gave perplexity above 400,000 on random tokens. These are smoke tests on a handful of sentences, not benchmarks.
+- **Versions.** The test suite passes on `transformers` 4.46.3, 4.57.3 and 5.22.2 with PyTorch 2.15.
+
+Not covered: the full checkpoint has not been run on a GPU or in bfloat16 through this code, and the continuations in the Prompting table were produced on CPU with weights rounded through float16. Multi-GPU `device_map` placement was only exercised on a small model.
+
+## Limitations
+
+- It is a base model with no instruction or safety tuning. It can produce incorrect, biased or offensive text, and it will state false things fluently.
+- Training data is English web text, maths and code. Expect weak results in other languages.
+- Sequences longer than 4,096 tokens were never seen in training. Up to 8,192 the model runs, but quality there has not been measured.
+- No evaluation results are available yet.
+
+## Fine-tuning for chat
+
+To continue with the fine-tuning stage that `train.py` defines, format each message as
+
+```
+<|{role}|>\n{content}\n<|endoftext|>
+```
+
+and compute the loss on assistant messages only, including their closing `<|endoftext|>`. The role markers are ordinary text, not special tokens. A generation prompt ends with `<|assistant|>\n`.

config.json

@@ -1,4 +1,13 @@
 {
+  "architectures": [
+    "CometForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_comet.CometConfig",
+    "AutoModel": "modeling_comet.CometModel",
+    "AutoModelForCausalLM": "modeling_comet.CometForCausalLM"
+  },
+  "model_type": "comet",
   "variant": "7.0b",
   "vocab_size": 50277,
   "hidden_size": 4096,
@@ -14,8 +23,15 @@
   "layer_pattern": ["global", "global", "global", "recall"],
   "layer_scale_init": 0.5,
   "qk_norm": true,
-  "tie_word_embeddings": true,
+  "rms_norm_eps": 1e-06,
   "fused_projections": true,
   "dropout": 0.0,
-  "gradient_checkpointing": true
-}
+  "initializer_range": 0.02,
+  "tie_word_embeddings": true,
+  "use_cache": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 1,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "dtype": "float32"
+}

configuration_comet.py

@@ -0,0 +1,145 @@
+"""Configuration for Comet-1 (OpenModel-1), a hybrid attention / causal-convolution language model."""
+
+from transformers import PretrainedConfig
+
+
+class CometConfig(PretrainedConfig):
+    r"""
+    Configuration class for [`CometModel`] and [`CometForCausalLM`].
+
+    The field names mirror the ``model_config`` block that the OpenModel-1 training
+    pipeline (``train.py``) writes next to every checkpoint, so a training-time
+    ``config.json`` loads unchanged.
+
+    Args:
+        vocab_size (`int`, defaults to 50277):
+            Size of the token embedding table (``len(tokenizer)`` of the GPT-NeoX-20B BPE tokenizer).
+        hidden_size (`int`, defaults to 4096):
+            Width of the residual stream.
+        num_hidden_layers (`int`, defaults to 36):
+            Number of blocks.
+        num_attention_heads (`int`, defaults to 32):
+            Number of query heads in the "global" (attention) blocks.
+        num_key_value_heads (`int`, defaults to 8):
+            Number of key/value heads (grouped-query attention).
+        intermediate_size (`int`, defaults to 12288):
+            Inner width of the SwiGLU feed-forward.
+        max_position_embeddings (`int`, defaults to 8192):
+            Number of positions in the precomputed RoPE table. Inputs longer than this are rejected.
+        rope_theta (`float`, defaults to 1000000.0):
+            RoPE base frequency.
+        memory_width (`int`, defaults to 2048):
+            Channel count of the depthwise convolution lane in "recall" blocks.
+        memory_kernel_size (`int`, defaults to 7):
+            Kernel size of the depthwise causal convolution.
+        memory_dilations (`list[int]`, defaults to `[1, 2, 4, 8]`):
+            Dilation table. Block ``i`` uses ``memory_dilations[i % len(memory_dilations)]``.
+        layer_pattern (`list[str]`, defaults to `["global", "global", "global", "recall"]`):
+            Repeating block schedule. Block ``i`` is of kind ``layer_pattern[i % len(layer_pattern)]``;
+            ``"global"`` is grouped-query attention, anything else is the causal recall mixer.
+        layer_scale_init (`float`, defaults to 0.5):
+            Initial value of the learned per-channel residual scales.
+        qk_norm (`bool`, defaults to `True`):
+            Apply RMSNorm to queries and keys (per head) before RoPE.
+        rms_norm_eps (`float`, defaults to 1e-6):
+            Epsilon of every RMSNorm.
+        fused_projections (`bool`, defaults to `True`):
+            Parameter layout. ``True`` stores Q/K/V as one ``qkv_proj`` matrix and the two SwiGLU input
+            projections as one ``gate_up`` matrix (the layout of the released checkpoint).
+        dropout (`float`, defaults to 0.0):
+            Dropout on both residual branches (training only).
+        initializer_range (`float`, defaults to 0.02):
+            Standard deviation used when initialising new weights.
+        tie_word_embeddings (`bool`, defaults to `True`):
+            Share the input embedding matrix with the output projection.
+        use_cache (`bool`, defaults to `True`):
+            Return a key/value + convolution-state cache from the forward pass.
+        variant (`str`, *optional*):
+            Free-form size label written by the training pipeline (for example ``"7.0b"``).
+    """
+
+    model_type = "comet"
+
+    def __init__(
+        self,
+        vocab_size=50277,
+        hidden_size=4096,
+        num_hidden_layers=36,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        intermediate_size=12288,
+        max_position_embeddings=8192,
+        rope_theta=1_000_000.0,
+        memory_width=2048,
+        memory_kernel_size=7,
+        memory_dilations=(1, 2, 4, 8),
+        layer_pattern=("global", "global", "global", "recall"),
+        layer_scale_init=0.5,
+        qk_norm=True,
+        rms_norm_eps=1e-6,
+        fused_projections=True,
+        dropout=0.0,
+        initializer_range=0.02,
+        tie_word_embeddings=True,
+        use_cache=True,
+        variant=None,
+        bos_token_id=0,
+        eos_token_id=0,
+        pad_token_id=1,
+        **kwargs,
+    ):
+        # `gradient_checkpointing` is a training-time switch that the training pipeline
+        # stores in its config dump. It is not an architecture property; enable it with
+        # `model.gradient_checkpointing_enable()` instead.
+        kwargs.pop("gradient_checkpointing", None)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = float(rope_theta)
+        self.memory_width = memory_width
+        self.memory_kernel_size = memory_kernel_size
+        self.memory_dilations = list(memory_dilations)
+        self.layer_pattern = list(layer_pattern)
+        self.layer_scale_init = layer_scale_init
+        self.qk_norm = qk_norm
+        self.rms_norm_eps = rms_norm_eps
+        self.fused_projections = fused_projections
+        self.dropout = dropout
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.variant = variant
+
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+        if (hidden_size // num_attention_heads) % 2 != 0:
+            raise ValueError("the attention head dimension must be even for RoPE")
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if memory_kernel_size < 1:
+            raise ValueError("memory_kernel_size must be positive")
+        if len(self.layer_pattern) == 0 or len(self.memory_dilations) == 0:
+            raise ValueError("layer_pattern and memory_dilations must not be empty")
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def block_kind(self, index):
+        """Kind of block ``index``: ``"global"`` (attention) or ``"recall"`` (causal convolution)."""
+        return self.layer_pattern[index % len(self.layer_pattern)]
+
+    def block_dilation(self, index):
+        """Convolution dilation that block ``index`` uses if it is a recall block."""
+        return self.memory_dilations[index % len(self.memory_dilations)]
+
+
+__all__ = ["CometConfig"]

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 1
+}

modeling_comet.py

@@ -0,0 +1,572 @@
+"""
+Comet-1 (OpenModel-1) for Hugging Face Transformers.
+
+Inference/fine-tuning port of the ``OpenModel1`` module defined in the training
+pipeline (``train.py``). Parameter names and shapes are identical to the training
+checkpoint, so ``model.safetensors`` loads without any key remapping.
+
+Architecture
+------------
+A pre-norm decoder with tied input/output embeddings. Blocks follow a repeating
+schedule (``config.layer_pattern``, by default ``[global, global, global, recall]``):
+
+* ``global`` blocks use causal grouped-query attention with per-head RMSNorm on
+  queries and keys (QK-norm) followed by rotary position embeddings.
+* ``recall`` blocks replace attention with a gated, depthwise, dilated *causal*
+  1-D convolution over a narrower lane (``config.memory_width``).
+
+Every block then applies a SwiGLU feed-forward. Both residual branches are scaled
+by a learned per-channel vector (LayerScale).
+
+Compared with ``train.py`` this file adds what generation needs and the training
+loop did not: an incremental cache (attention keys/values *and* convolution
+history), padding-aware attention masks, and the standard Transformers outputs.
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers import DynamicCache, GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .configuration_comet import CometConfig
+
+
+def _cache_length(cache) -> int:
+    """Number of tokens already stored in ``cache`` (0 for ``None`` or an empty cache)."""
+    if cache is None:
+        return 0
+    try:
+        return int(cache.get_seq_length())
+    except Exception:
+        return 0
+
+
+class CometRMSNorm(nn.Module):
+    def __init__(self, size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        weight = self.weight.to(x.dtype)
+        if hasattr(F, "rms_norm"):
+            return F.rms_norm(x, (x.shape[-1],), weight, self.eps)
+        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
+        return weight * (x * torch.rsqrt(variance + self.eps).to(x.dtype))
+
+    def extra_repr(self) -> str:
+        return f"{self.weight.shape[0]}, eps={self.eps}"
+
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Rotate interleaved (even, odd) channel pairs. ``x``: [B, H, T, D]; ``cos``/``sin``: [B|1, 1, T, D/2]."""
+    even = x[..., ::2]
+    odd = x[..., 1::2]
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+class CometAttention(nn.Module):
+    """Causal grouped-query attention with QK-norm and RoPE."""
+
+    def __init__(self, config: CometConfig):
+        super().__init__()
+        self.heads = config.num_attention_heads
+        self.kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.q_dim = self.heads * self.head_dim
+        self.kv_dim = self.kv_heads * self.head_dim
+        self.fused_projections = config.fused_projections
+        if self.fused_projections:
+            self.qkv_proj = nn.Linear(config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=False)
+        else:
+            self.q_proj = nn.Linear(config.hidden_size, self.q_dim, bias=False)
+            self.k_proj = nn.Linear(config.hidden_size, self.kv_dim, bias=False)
+            self.v_proj = nn.Linear(config.hidden_size, self.kv_dim, bias=False)
+        self.o_proj = nn.Linear(self.q_dim, config.hidden_size, bias=False)
+        self.o_proj._comet_residual_out = True
+        if config.qk_norm:
+            self.q_norm = CometRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = CometRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        else:
+            self.q_norm = nn.Identity()
+            self.k_norm = nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+        cache=None,
+        layer_idx: int = 0,
+    ) -> torch.Tensor:
+        batch, sequence, _ = x.shape
+        if self.fused_projections:
+            q, k, v = self.qkv_proj(x).split((self.q_dim, self.kv_dim, self.kv_dim), dim=-1)
+        else:
+            q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        q = q.view(batch, sequence, self.heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
+
+        cos = cos.to(device=q.device, dtype=q.dtype)
+        sin = sin.to(device=q.device, dtype=q.dtype)
+        q = _apply_rope(self.q_norm(q), cos, sin)
+        k = _apply_rope(self.k_norm(k), cos, sin)
+
+        if cache is not None:
+            # Keys are stored already normalised and rotated, at key/value-head granularity.
+            k, v = cache.update(k, v, layer_idx)
+
+        groups = self.heads // self.kv_heads
+        if groups > 1:
+            # Query head h reads key/value head h // groups (same grouping as `enable_gqa`).
+            length = k.shape[-2]
+            k = k[:, :, None].expand(batch, self.kv_heads, groups, length, self.head_dim)
+            k = k.reshape(batch, self.heads, length, self.head_dim)
+            v = v[:, :, None].expand(batch, self.kv_heads, groups, length, self.head_dim)
+            v = v.reshape(batch, self.heads, length, self.head_dim)
+
+        if attn_mask is not None:
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.to(q.device))
+        else:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
+        return self.o_proj(y.transpose(1, 2).reshape(batch, sequence, self.q_dim))
+
+
+class CometRecallMixer(nn.Module):
+    """Gated local recall branch: depthwise, dilated, causal convolution over a narrow lane."""
+
+    def __init__(self, config: CometConfig, dilation: int):
+        super().__init__()
+        width = config.memory_width
+        self.width = width
+        self.kernel_size = config.memory_kernel_size
+        self.dilation = dilation
+        # Number of past frames a single output position can see.
+        self.history = (self.kernel_size - 1) * self.dilation
+        self.in_proj = nn.Linear(config.hidden_size, 2 * width, bias=False)
+        self.depthwise = nn.Conv1d(
+            width,
+            width,
+            kernel_size=self.kernel_size,
+            dilation=dilation,
+            groups=width,
+            bias=False,
+        )
+        self.out_proj = nn.Linear(width, config.hidden_size, bias=False)
+        self.out_proj._comet_residual_out = True
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        token_mask: Optional[torch.Tensor] = None,
+        cache=None,
+        layer_idx: int = 0,
+    ) -> torch.Tensor:
+        batch, sequence, _ = x.shape
+        gate, value = self.in_proj(x).chunk(2, dim=-1)
+        if token_mask is not None:
+            # Padding tokens contribute nothing to the convolution, exactly like the
+            # zero padding that precedes the first real token during training.
+            value = value * token_mask.to(device=value.device, dtype=value.dtype).unsqueeze(-1)
+
+        if cache is not None:
+            # The convolution input is cached in the "key" slot as [B, 1, T, width] so
+            # that every generic cache operation (beam reordering, cropping, batch
+            # expansion) applies to it unchanged. The "value" slot is a placeholder.
+            frames, _ = cache.update(value.unsqueeze(1), value.new_zeros(batch, 1, sequence, 1), layer_idx)
+            context = frames[:, 0, -(self.history + sequence) :]
+        else:
+            context = value
+
+        left = self.history + sequence - context.shape[1]
+        context = F.pad(context.transpose(1, 2), (left, 0))
+        value = self.depthwise(context).transpose(1, 2)
+        return self.out_proj(torch.sigmoid(gate) * value)
+
+
+class CometSwiGLU(nn.Module):
+    def __init__(self, config: CometConfig):
+        super().__init__()
+        self.fused_projections = config.fused_projections
+        if self.fused_projections:
+            self.gate_up = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+        else:
+            self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+            self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.down._comet_residual_out = True
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.fused_projections:
+            gate, up = self.gate_up(x).chunk(2, dim=-1)
+        else:
+            gate, up = self.gate(x), self.up(x)
+        return self.down(F.silu(gate) * up)
+
+
+class CometBlock(nn.Module):
+    def __init__(self, config: CometConfig, index: int):
+        super().__init__()
+        self.index = index
+        self.kind = config.block_kind(index)
+        self.norm1 = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm2 = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        if self.kind == "global":
+            self.mixer = CometAttention(config)
+        else:
+            self.mixer = CometRecallMixer(config, config.block_dilation(index))
+        self.mlp = CometSwiGLU(config)
+        self.mixer_scale = nn.Parameter(torch.full((config.hidden_size,), float(config.layer_scale_init)))
+        self.mlp_scale = nn.Parameter(torch.full((config.hidden_size,), float(config.layer_scale_init)))
+        self.dropout = config.dropout
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        token_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+        cache=None,
+    ) -> torch.Tensor:
+        h = self.norm1(x)
+        if self.kind == "global":
+            h = self.mixer(h, cos, sin, attn_mask=attn_mask, is_causal=is_causal, cache=cache, layer_idx=self.index)
+        else:
+            h = self.mixer(h, token_mask=token_mask, cache=cache, layer_idx=self.index)
+        if self.dropout:
+            h = F.dropout(h, p=self.dropout, training=self.training)
+        x = x + h * self.mixer_scale.to(h.dtype)
+        h = self.mlp(self.norm2(x))
+        if self.dropout:
+            h = F.dropout(h, p=self.dropout, training=self.training)
+        return x + h * self.mlp_scale.to(h.dtype)
+
+
+class CometPreTrainedModel(PreTrainedModel):
+    config_class = CometConfig
+    base_model_prefix = ""
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["CometBlock"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_cache_class = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            if getattr(module, "_comet_residual_out", False):
+                std = std / math.sqrt(2.0 * self.config.num_hidden_layers)
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, (nn.Embedding, nn.Conv1d)):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, CometRMSNorm):
+            nn.init.ones_(module.weight)
+        elif isinstance(module, CometBlock):
+            nn.init.constant_(module.mixer_scale, float(self.config.layer_scale_init))
+            nn.init.constant_(module.mlp_scale, float(self.config.layer_scale_init))
+
+
+class CometModel(CometPreTrainedModel):
+    """Bare Comet trunk: token embeddings, blocks and the final norm. Returns hidden states."""
+
+    def __init__(self, config: CometConfig):
+        super().__init__(config)
+        self._build_trunk(config)
+        self.post_init()
+
+    def _build_trunk(self, config: CometConfig) -> None:
+        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.blocks = nn.ModuleList(CometBlock(config, index) for index in range(config.num_hidden_layers))
+        self.norm = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.gradient_checkpointing = False
+
+        # RoPE tables are built lazily on first use and kept as plain attributes rather than
+        # buffers: they are derived constants, and this keeps them out of every loading path
+        # (meta-device instantiation, device maps, dtype casts).
+        self._rope_tables = None
+
+    def _rope(self, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+        tables = self._rope_tables
+        if tables is None or tables[0].device != device:
+            config = self.config
+            head_dim = config.hidden_size // config.num_attention_heads
+            exponent = torch.arange(0, head_dim, 2, dtype=torch.float32, device="cpu") / head_dim
+            inv_freq = 1.0 / (config.rope_theta**exponent)
+            positions = torch.arange(config.max_position_embeddings, dtype=torch.float32, device="cpu")
+            frequencies = torch.outer(positions, inv_freq)
+            tables = (frequencies.cos().to(device), frequencies.sin().to(device))
+            self._rope_tables = tables
+        return tables
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed = value
+
+    def _run_trunk(
+        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.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ):
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of `input_ids` or `inputs_embeds`.")
+        if inputs_embeds is None:
+            inputs_embeds = self.embed(input_ids)
+        batch, sequence = inputs_embeds.shape[:2]
+        device = inputs_embeds.device
+
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+        if self.training and past_key_values is None:
+            use_cache = False
+        cache = past_key_values
+        if cache is not None and not hasattr(cache, "update"):
+            raise TypeError(
+                "`past_key_values` must be a `transformers.Cache` instance (for example `DynamicCache()`); "
+                "legacy tuple caches are not supported."
+            )
+        if cache is None and use_cache:
+            cache = DynamicCache()
+        past = _cache_length(cache)
+        total = past + sequence
+        if total > self.config.max_position_embeddings:
+            raise ValueError(
+                f"Sequence length {total} exceeds max_position_embeddings={self.config.max_position_embeddings}."
+            )
+
+        # ---- masks -------------------------------------------------------------
+        attn_mask = None  # [B|1, 1, T, past + T] boolean, True = may attend
+        token_mask = None  # [B, T] boolean, True = real token (current chunk only)
+        is_causal = False
+        keep = None
+        if attention_mask is not None:
+            if attention_mask.dim() != 2:
+                raise ValueError("`attention_mask` must have shape [batch, sequence].")
+            keep = attention_mask.to(device=device, dtype=torch.bool)
+            if keep.shape[1] == sequence and past > 0:
+                # Only the new chunk was described; everything cached is treated as real.
+                keep = torch.cat([keep.new_ones(batch, past), keep], dim=1)
+            if keep.shape[1] != total:
+                raise ValueError(
+                    f"`attention_mask` has length {attention_mask.shape[1]} but the model sees {total} positions "
+                    f"({past} cached + {sequence} new)."
+                )
+            if bool(keep.all()):
+                keep = None
+
+        if keep is not None or (past > 0 and sequence > 1):
+            query_positions = past + torch.arange(sequence, device=device)
+            key_positions = torch.arange(total, device=device)
+            attn_mask = (key_positions[None, :] <= query_positions[:, None])[None, None]
+            if keep is not None:
+                # Every position may always attend to itself so that padded rows never
+                # become fully masked (which would produce NaNs).
+                diagonal = (key_positions[None, :] == query_positions[:, None])[None, None]
+                attn_mask = (attn_mask & keep[:, None, None, :]) | diagonal
+                token_mask = keep[:, past:]
+        elif past == 0 and sequence > 1:
+            is_causal = True
+
+        # ---- positions ---------------------------------------------------------
+        if position_ids is None:
+            if keep is not None:
+                # Real tokens are numbered 0, 1, 2, ... regardless of how much padding precedes them.
+                position_ids = (keep.long().cumsum(-1) - 1).clamp_min(0)[:, past:]
+            else:
+                position_ids = torch.arange(past, total, device=device)[None, :]
+        else:
+            position_ids = position_ids.to(device)
+            if position_ids.dim() == 1:
+                position_ids = position_ids[None, :]
+            if position_ids.shape[-1] != sequence:
+                position_ids = position_ids[..., -sequence:]
+            if int(position_ids.max()) >= self.config.max_position_embeddings:
+                raise ValueError("`position_ids` exceed max_position_embeddings.")
+        rope_cos, rope_sin = self._rope(device)
+        cos = rope_cos[position_ids].unsqueeze(1)
+        sin = rope_sin[position_ids].unsqueeze(1)
+
+        # ---- blocks ------------------------------------------------------------
+        x = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
+        for block in self.blocks:
+            if output_hidden_states:
+                all_hidden_states += (x,)
+            if self.gradient_checkpointing and self.training and cache is None:
+                x = torch.utils.checkpoint.checkpoint(
+                    block, x, cos, sin, attn_mask, token_mask, is_causal, use_reentrant=False
+                )
+            else:
+                x = block(x, cos, sin, attn_mask=attn_mask, token_mask=token_mask, is_causal=is_causal, cache=cache)
+        x = self.norm(x)
+        if output_hidden_states:
+            all_hidden_states += (x,)
+        return x, (cache if use_cache else None), all_hidden_states
+
+    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.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        if output_hidden_states is None:
+            output_hidden_states = bool(getattr(self.config, "output_hidden_states", False))
+        if return_dict is None:
+            return_dict = bool(getattr(self.config, "return_dict", True))
+        hidden, cache, all_hidden_states = self._run_trunk(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+        )
+        if not return_dict:
+            return tuple(item for item in (hidden, cache, all_hidden_states) if item is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden,
+            past_key_values=cache,
+            hidden_states=all_hidden_states,
+        )
+
+
+class CometForCausalLM(CometModel, GenerationMixin):
+    """
+    Comet trunk with the language-modelling head.
+
+    The state-dict layout is flat and identical to the training checkpoint
+    (``embed.weight``, ``blocks.N.*``, ``norm.weight``). With
+    ``config.tie_word_embeddings=True`` there is no separate head parameter: logits are
+    computed directly against ``embed.weight``.
+    """
+
+    def __init__(self, config: CometConfig):
+        CometPreTrainedModel.__init__(self, config)
+        self._build_trunk(config)
+        if config.tie_word_embeddings:
+            self.lm_head = None
+        else:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self):
+        # `None` when tied: there is nothing for Transformers to tie or resize separately.
+        return self.lm_head
+
+    def set_output_embeddings(self, value) -> None:
+        self.lm_head = value
+
+    def _logits(self, hidden: torch.Tensor) -> torch.Tensor:
+        weight = self.embed.weight if self.lm_head is None else self.lm_head.weight
+        return F.linear(hidden.to(weight.device), weight)
+
+    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.Tensor] = 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,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch, sequence)`, *optional*):
+            Next-token targets. They are shifted inside the model, so pass ``labels=input_ids``;
+            positions set to ``-100`` are ignored.
+        """
+        if output_hidden_states is None:
+            output_hidden_states = bool(getattr(self.config, "output_hidden_states", False))
+        if return_dict is None:
+            return_dict = bool(getattr(self.config, "return_dict", True))
+        hidden, cache, all_hidden_states = self._run_trunk(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+        )
+        logits = self._logits(hidden)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1].float()
+            shift_labels = labels[:, 1:].to(shift_logits.device)
+            total = F.cross_entropy(
+                shift_logits.reshape(-1, shift_logits.shape[-1]),
+                shift_labels.reshape(-1),
+                ignore_index=-100,
+                reduction="sum",
+            )
+            loss = total / shift_labels.ne(-100).sum().clamp_min(1)
+
+        if not return_dict:
+            return tuple(item for item in (loss, logits, cache, all_hidden_states) if item is not None)
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=cache,
+            hidden_states=all_hidden_states,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values=None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        """Feed only the tokens the cache has not seen yet; keep the full-length attention mask."""
+        past = _cache_length(past_key_values)
+        model_inputs = {}
+        if inputs_embeds is not None and past == 0:
+            model_inputs["inputs_embeds"] = inputs_embeds
+        else:
+            seen = past
+            if inputs_embeds is not None:
+                # Generation started from embeddings: `input_ids` holds only the generated tokens.
+                seen = past - inputs_embeds.shape[1]
+            unseen = input_ids.shape[1] - seen
+            if past > 0:
+                input_ids = input_ids[:, -unseen:] if unseen > 0 else input_ids[:, -1:]
+            model_inputs["input_ids"] = input_ids.contiguous()
+        model_inputs["past_key_values"] = past_key_values
+        model_inputs["attention_mask"] = attention_mask
+        model_inputs["use_cache"] = kwargs.get("use_cache", True)
+        if kwargs.get("output_hidden_states") is not None:
+            model_inputs["output_hidden_states"] = kwargs["output_hidden_states"]
+        return model_inputs
+
+
+__all__ = ["CometPreTrainedModel", "CometModel", "CometForCausalLM"]

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>",
+  "pad_token": "<|padding|>"
+}

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>",
+  "pad_token": "<|padding|>",
+  "add_prefix_space": false,
+  "clean_up_tokenization_spaces": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "padding_side": "left"
+}