Upload repaired SpliNet 2B-token pretrained model

README.md

@@ -0,0 +1,50 @@
+---
+language:
+- en
+pipeline_tag: fill-mask
+library_name: transformers
+datasets:
+- allenai/c4
+---
+
+# SpliNet
+
+**SpliNet: A Zero-Parameter B-Spline Transformer with Linear Complexity**
+
+SpliNet replaces learned self-attention token mixing with a fixed order-2 single-sided cardinal B-spline operator.
+
+This model was pretrained from scratch on exactly **2,000,000,000 C4 tokens** using the dedicated SpliNet tokenizer.
+
+## Architecture
+
+- Layers: 12
+- Hidden size: 768
+- Heads: 12
+- FFN width: 3072
+- Sequence length: 512
+- Vocabulary: 32000
+- Spline order: 2
+- Spline radius: 16
+- Trainable mixer parameters: 0
+- Total parameters: 82,894,592
+- Trainable parameters: 82,894,592
+
+## Pretraining
+
+- Training tokens: 2,000,000,000
+- Validation tokens: 5,120,000
+- Objective: masked language modeling
+- Masked positions: 77/512
+- Optimizer: AdamW
+- Precision: BF16
+- Hardware: NVIDIA A100-SXM4-80GB
+
+## Final validation
+
+- MLM loss: 4.614502
+- MLM perplexity: 100.937575
+
+Load with `trust_remote_code=True`.
+
+OpenReview: https://openreview.net/forum?id=nWHnuiEF3C
+GitHub: https://github.com/AngshulMajumdar/SpliNet

checksums.sha256

@@ -0,0 +1,13 @@
+2d50d51e32c691a1034ed87770298f054aaf7fcf6e3046274a3b528f3d0135d0  README.md
+e22e250a1383a774e2842c1f9183895ea7678b6974114489aec8752b06f4dcb8  config.json
+91c0167660e96c39d55a31053c58f87abbcbe58cf22525dfb9b90ce6d28bdd3e  configuration_splinet.py
+a891f4a978c2551fc8851346601fa3d560701ba69471272c4c5a39f4d70fbdd6  model.safetensors
+68b78991d599f1c6c3b52d3882bb07e65c81243500064b9317e0f289f57576ba  modeling_splinet.py
+321813b734e7c50df348e53600d5b90ba9d204eba535545279babac85babe2ab  requirements.txt
+ace55eb0f41e170e8f3c891dd4ab707a2d871299ab7e4f4116c6a9d3128f31cf  special_tokens_map.json
+119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d  spiece.model
+ec799958399f5fcd35f12ac451e6b7517e2eda01f50a7b9d2089c1425e9a0d05  spiece.vocab
+9e0b42c275c87f712d095f002bbfa5a435ed89250785ef9b4d4ee9a7ab18a53e  tokenization_splinet.py
+8171bc64afff9fe2f89ac1e011d0b829d4b5413b398ef7b8b9f408ad5cd30325  tokenizer_config.json
+d80f0c668bf183e6f327c167bb2723f0e565c946c05d6b5587afec7ca10575fb  tokenizer_metadata.json
+8f64420de4f4a6cd9ec7c8920f8be32463c8ba32d85dc9ab86865c3aba01e804  training_metadata.json

config.json

@@ -0,0 +1,36 @@
+{
+  "architectures": [
+    "SpliNetForMaskedLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_splinet.SpliNetConfig",
+    "AutoModel": "modeling_splinet.SpliNetModel",
+    "AutoModelForMaskedLM": "modeling_splinet.SpliNetForMaskedLM"
+  },
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "hidden_act": "gelu_new",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "splinet",
+  "num_hidden_layers": 12,
+  "pad_token_id": 3,
+  "splinet_num_heads": 12,
+  "splinet_order": 2,
+  "splinet_radius": 16,
+  "splinet_sidedness": "single",
+  "tie_word_embeddings": true,
+  "tokenizer_class": "SpliNetTokenizer",
+  "tokenizer_sha256": "119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d",
+  "tpu_short_seq_length": 512,
+  "training_tokens": 2000000000,
+  "transformers_version": "5.16.1",
+  "type_vocab_size": 4,
+  "use_tpu_fourier_optimizations": false,
+  "vocab_size": 32000
+}

configuration_splinet.py

@@ -0,0 +1,24 @@
+from transformers import FNetConfig
+
+
+class SpliNetConfig(FNetConfig):
+    model_type = "splinet"
+
+    def __init__(
+        self,
+        splinet_num_heads=12,
+        splinet_radius=16,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.splinet_num_heads = int(
+            splinet_num_heads
+        )
+
+        self.splinet_radius = int(
+            splinet_radius
+        )
+
+        self.splinet_order = 2
+        self.splinet_sidedness = "single"

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a891f4a978c2551fc8851346601fa3d560701ba69471272c4c5a39f4d70fbdd6
+size 331590896

modeling_splinet.py

@@ -0,0 +1,366 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+)
+
+from transformers.models.fnet.modeling_fnet import (
+    FNetEmbeddings,
+    FNetIntermediate,
+    FNetOnlyMLMHead,
+    FNetOutput,
+    FNetPreTrainedModel,
+)
+
+from .configuration_splinet import SpliNetConfig
+
+
+class SplineMixer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.radius = int(
+            config.splinet_radius
+        )
+
+        z = -3.0 + 2.0 * math.sqrt(2.0)
+
+        # Fixed analytic spline coefficients.
+        #
+        # IMPORTANT:
+        # Do NOT register these as a non-persistent tensor buffer.
+        # Hugging Face low-memory/meta-device loading can materialize
+        # such a buffer without its analytically initialized values.
+        #
+        # Store the 33 coefficients as ordinary Python floats instead.
+        # They are recreated on the actual input device and dtype in
+        # forward(). They are not learned model state.
+        self.kernel_values = tuple(
+            float(
+                math.sqrt(2.0)
+                * (z ** abs(k))
+            )
+            for k in range(
+                -self.radius,
+                self.radius + 1,
+            )
+        )
+
+    def forward(self, x):
+        if x.shape[1] != 512:
+            raise ValueError(
+                "This released SpliNet checkpoint was pretrained "
+                "and validated on fixed 512-token blocks. "
+                f"Received sequence length {x.shape[1]}. "
+                "Tokenize/pack the input to exactly 512 tokens."
+            )
+
+        d = x.shape[-1]
+
+        y = (
+            x.transpose(1, 2)
+            .contiguous()
+        )
+
+        y = F.pad(
+            y,
+            (
+                self.radius,
+                self.radius,
+            ),
+            mode="reflect",
+        )
+
+        kernel = (
+            torch.tensor(
+                self.kernel_values,
+                device=y.device,
+                dtype=y.dtype,
+            )
+            .view(1, 1, -1)
+            .expand(d, 1, -1)
+            .contiguous()
+        )
+
+        y = F.conv1d(
+            y,
+            kernel,
+            groups=d,
+        )
+
+        return (
+            y.transpose(1, 2)
+            .contiguous()
+        )
+
+
+class SpliNetMixingBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.mixer = SplineMixer(
+            config
+        )
+
+        self.LayerNorm = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+        )
+
+    def forward(self, x):
+        return self.LayerNorm(
+            x + self.mixer(x)
+        )
+
+
+class SpliNetLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.mixing = (
+            SpliNetMixingBlock(
+                config
+            )
+        )
+
+        self.intermediate = (
+            FNetIntermediate(
+                config
+            )
+        )
+
+        self.output = (
+            FNetOutput(
+                config
+            )
+        )
+
+    def forward(self, x):
+        x = self.mixing(x)
+
+        return self.output(
+            self.intermediate(x),
+            x,
+        )
+
+
+class SpliNetEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.layer = nn.ModuleList(
+            [
+                SpliNetLayer(config)
+                for _ in range(
+                    config.num_hidden_layers
+                )
+            ]
+        )
+
+    def forward(
+        self,
+        x,
+        output_hidden_states=False,
+    ):
+        hidden_states = (
+            ()
+            if output_hidden_states
+            else None
+        )
+
+        for layer in self.layer:
+            if output_hidden_states:
+                hidden_states += (x,)
+
+            x = layer(x)
+
+        if output_hidden_states:
+            hidden_states += (x,)
+
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states,
+        )
+
+
+class SpliNetModel(FNetPreTrainedModel):
+    config_class = SpliNetConfig
+    base_model_prefix = "splinet"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = (
+            FNetEmbeddings(config)
+        )
+
+        self.encoder = (
+            SpliNetEncoder(config)
+        )
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return (
+            self.embeddings
+            .word_embeddings
+        )
+
+    def set_input_embeddings(
+        self,
+        value,
+    ):
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        output_hidden_states=False,
+        **kwargs,
+    ):
+        if input_ids is not None:
+            shape = input_ids.shape
+            device = input_ids.device
+
+        elif inputs_embeds is not None:
+            shape = inputs_embeds.shape[:-1]
+            device = inputs_embeds.device
+
+        else:
+            raise ValueError(
+                "input_ids or inputs_embeds required."
+            )
+
+        if shape[1] != 512:
+            raise ValueError(
+                "This SpliNet checkpoint requires exactly "
+                f"512 tokens; received {shape[1]}."
+            )
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(
+                shape,
+                dtype=torch.long,
+                device=device,
+            )
+
+        x = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+        )
+
+        return self.encoder(
+            x,
+            output_hidden_states=output_hidden_states,
+        )
+
+
+class SpliNetForMaskedLM(FNetPreTrainedModel):
+    config_class = SpliNetConfig
+    base_model_prefix = "splinet"
+
+    _tied_weights_keys = {
+        "cls.predictions.decoder.bias":
+            "cls.predictions.bias",
+
+        "cls.predictions.decoder.weight":
+            "splinet.embeddings.word_embeddings.weight",
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.splinet = (
+            SpliNetModel(config)
+        )
+
+        self.cls = (
+            FNetOnlyMLMHead(config)
+        )
+
+        self.post_init()
+
+        if config.tie_word_embeddings:
+            (
+                self.cls
+                .predictions
+                .decoder.weight
+            ) = (
+                self.splinet
+                .embeddings
+                .word_embeddings
+                .weight
+            )
+
+    def get_input_embeddings(self):
+        return (
+            self.splinet
+            .embeddings
+            .word_embeddings
+        )
+
+    def set_input_embeddings(
+        self,
+        value,
+    ):
+        (
+            self.splinet
+            .embeddings
+            .word_embeddings
+        ) = value
+
+    def get_output_embeddings(self):
+        return (
+            self.cls
+            .predictions
+            .decoder
+        )
+
+    def set_output_embeddings(
+        self,
+        value,
+    ):
+        self.cls.predictions.decoder = value
+
+    def forward(
+        self,
+        input_ids=None,
+        labels=None,
+        **kwargs,
+    ):
+        out = self.splinet(
+            input_ids=input_ids,
+            **kwargs,
+        )
+
+        logits = self.cls(
+            out.last_hidden_state
+        )
+
+        loss = None
+
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.reshape(
+                    -1,
+                    self.config.vocab_size,
+                ),
+                labels.reshape(-1),
+                ignore_index=-100,
+            )
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=out.hidden_states,
+        )

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers>=4.45.0
+sentencepiece>=0.2.0
+safetensors>=0.4.5

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "unk_token": "<unk>",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "cls_token": "<cls>",
+  "sep_token": "<sep>",
+  "mask_token": "<mask>"
+}

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d
+size 806898

tokenization_splinet.py

@@ -0,0 +1,153 @@
+import os
+import shutil
+import unicodedata
+import sentencepiece as spm
+
+from transformers import PreTrainedTokenizer
+
+
+class SpliNetTokenizer(PreTrainedTokenizer):
+
+    vocab_files_names = {
+        "vocab_file": "spiece.model"
+    }
+
+    model_input_names = [
+        "input_ids",
+        "token_type_ids",
+        "attention_mask",
+    ]
+
+    def __init__(
+        self,
+        vocab_file,
+        do_lower_case=True,
+        **kwargs,
+    ):
+        self.vocab_file = vocab_file
+        self.do_lower_case = bool(do_lower_case)
+
+        self.sp_model = spm.SentencePieceProcessor(
+            model_file=vocab_file
+        )
+
+        # tokenizer_config.json may already provide these.
+        # setdefault prevents passing any keyword twice.
+        kwargs.setdefault("unk_token", "<unk>")
+        kwargs.setdefault("bos_token", "<s>")
+        kwargs.setdefault("eos_token", "</s>")
+        kwargs.setdefault("pad_token", "<pad>")
+        kwargs.setdefault("cls_token", "<cls>")
+        kwargs.setdefault("sep_token", "<sep>")
+        kwargs.setdefault("mask_token", "<mask>")
+
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self):
+        return int(
+            self.sp_model.get_piece_size()
+        )
+
+    def get_vocab(self):
+        return {
+            self.sp_model.id_to_piece(i): i
+            for i in range(self.vocab_size)
+        }
+
+    def _normalize(self, text):
+        text = text or ""
+
+        if self.do_lower_case:
+            text = unicodedata.normalize(
+                "NFKC",
+                text,
+            ).lower()
+
+        return " ".join(text.split())
+
+    def _tokenize(self, text):
+        return self.sp_model.encode(
+            self._normalize(text),
+            out_type=str,
+        )
+
+    def _convert_token_to_id(self, token):
+        return int(
+            self.sp_model.piece_to_id(token)
+        )
+
+    def _convert_id_to_token(self, index):
+        return self.sp_model.id_to_piece(
+            int(index)
+        )
+
+    def convert_tokens_to_string(self, tokens):
+        return self.sp_model.decode(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0,
+        token_ids_1=None,
+    ):
+        if token_ids_1 is None:
+            return (
+                [self.cls_token_id]
+                + list(token_ids_0)
+                + [self.sep_token_id]
+            )
+
+        return (
+            [self.cls_token_id]
+            + list(token_ids_0)
+            + [self.sep_token_id]
+            + list(token_ids_1)
+            + [self.sep_token_id]
+        )
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0,
+        token_ids_1=None,
+    ):
+        if token_ids_1 is None:
+            return [0] * (
+                len(token_ids_0) + 2
+            )
+
+        return (
+            [0] * (len(token_ids_0) + 2)
+            + [1] * (len(token_ids_1) + 1)
+        )
+
+    def save_vocabulary(
+        self,
+        save_directory,
+        filename_prefix=None,
+    ):
+        os.makedirs(
+            save_directory,
+            exist_ok=True,
+        )
+
+        prefix = (
+            filename_prefix + "-"
+            if filename_prefix
+            else ""
+        )
+
+        destination = os.path.join(
+            save_directory,
+            prefix + "spiece.model",
+        )
+
+        if (
+            os.path.abspath(self.vocab_file)
+            != os.path.abspath(destination)
+        ):
+            shutil.copy2(
+                self.vocab_file,
+                destination,
+            )
+
+        return (destination,)

tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "tokenizer_class": "SpliNetTokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_splinet.SpliNetTokenizer",
+      null
+    ]
+  },
+  "model_max_length": 512,
+  "do_lower_case": true,
+  "unk_token": "<unk>",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "cls_token": "<cls>",
+  "sep_token": "<sep>",
+  "mask_token": "<mask>"
+}

tokenizer_metadata.json

@@ -0,0 +1,21 @@
+{
+  "source": "allenai/c4",
+  "config": "en",
+  "split": "train",
+  "model_type": "unigram",
+  "vocab_size": 32000,
+  "lowercase": true,
+  "normalization": "NFKC + SentencePiece nmt_nfkc",
+  "training_text_bytes": 536874807,
+  "training_documents": 247312,
+  "special_token_ids": {
+    "<unk>": 0,
+    "<s>": 1,
+    "</s>": 2,
+    "<pad>": 3,
+    "<cls>": 4,
+    "<sep>": 5,
+    "<mask>": 6
+  },
+  "sha256_spiece_model": "119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d"
+}

training_metadata.json

@@ -0,0 +1,28 @@
+{
+  "model": "SpliNet",
+  "dataset": "allenai/c4 en",
+  "training_tokens": 2000000000,
+  "validation_tokens": 5120000,
+  "tokenizer_sha256": "119ec6b2af9cbbc56f297bd606b69f79f6e3130a34ee56122ee813a58d15bb9d",
+  "vocab_size": 32000,
+  "sequence_length": 512,
+  "layers": 12,
+  "hidden_size": 768,
+  "heads": 12,
+  "ffn_size": 3072,
+  "spline_order": 2,
+  "spline_radius": 16,
+  "spline_sidedness": "single",
+  "physical_batch": 160,
+  "effective_batch": 1024,
+  "optimizer": "AdamW",
+  "base_learning_rate": 0.0001,
+  "weight_decay": 0.01,
+  "precision": "bfloat16",
+  "gpu": "NVIDIA A100-SXM4-80GB",
+  "final_validation_loss": 4.6145022583007815,
+  "final_validation_perplexity": 100.93757518673716,
+  "best_validation_loss": 4.6145022583007815,
+  "openreview": "https://openreview.net/forum?id=nWHnuiEF3C",
+  "github": "https://github.com/AngshulMajumdar/SpliNet"
+}