Repository: keras-team/keras-nlp
Pull request #1403 (2024-01-17 02:11:25)

Bloom tokenizer
The conversion script worked alright and the backbone and tokenizer generated the same output as hugging face.
here is a gist: https://colab.research.google.com/gist/abuelnasr0/a603fd2077b6cc579b51876b8cefecd9/bloom.ipynb

I have also changed the preset names. @mattdangerw please take a look at them and tell me If you want to change them.

regarding uploading the model variations into kaggle. I have an issue. first of all I can't upload it from my local machine because I have bad internet connection. and we have limited internet. my internet quota will end If I tried to upload or download models. I only work at colab or kaggle. 
so I tried to put a link to a google drive folder containing the preset model and it didn't work.
I also tried to add variation from a notebook output. but kaggle doesn't show the notebook that outputs the preset as option to me, although it's public and has an Output Data. the notebook link: https://www.kaggle.com/code/mohamedabuelnasr/bloom-model-saving

any suggestions to upload the preset ?


**UPDATE**
I managed to upload the preset correctly after zipping the "assets" directory in the Kaggle output and copying it's link and other files link from kaggle notebook output.
But it would be useful if someone could tell me why my notebook output doesn't appear in the "Notebook output" section while uploading the model.

abuelnasr0: I managed to upload the preset into kaggle by copying each file link from kaggle notebook output. but I can't put the "merges.txt" and "vocabulary.json" files into "assets/tokenizer" directory
abuelnasr0: The Preset test will fail until I have permission to make the model public at kaggle. I tried loading the preset in a kaggle notebook because I have access to the link and it worked fine.
mattdangerw: @abuelnasr0 can you share it with me? https://www.kaggle.com/matthewdwatson

Public sharing is not rolled out to Kaggle quite yet.

abuelnasr0: @mattdangerw I have shared the model with you now.
mattdangerw: @abuelnasr0 Thanks!

I think this is already to go. Just copied the model you shared over to the Keras team. If tests look green, I will merge this in!

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 30736594..7b898138 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -36,6 +36,7 @@ from keras_nlp.models.bert.bert_masked_lm_preprocessor import (
 from keras_nlp.models.bert.bert_preprocessor import BertPreprocessor
 from keras_nlp.models.bert.bert_tokenizer import BertTokenizer
 from keras_nlp.models.bloom.bloom_backbone import BloomBackbone
+from keras_nlp.models.bloom.bloom_tokenizer import BloomTokenizer
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_classifier import (
     DebertaV3Classifier,
diff --git a/keras_nlp/models/bloom/bloom_backbone.py b/keras_nlp/models/bloom/bloom_backbone.py
index e3d66998..2fd18ed7 100644
--- a/keras_nlp/models/bloom/bloom_backbone.py
+++ b/keras_nlp/models/bloom/bloom_backbone.py
@@ -11,11 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
+
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.bloom.bloom_decoder import BloomDecoder
+from keras_nlp.models.bloom.bloom_presets import backbone_presets
+from keras_nlp.utils.python_utils import classproperty
 
 
 def _bloom_kernel_initializer(stddev=0.02):
@@ -35,7 +39,8 @@ class BloomBackbone(Backbone):
     load preset architectures and weights, use the `from_preset()` constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
-    warranties or conditions of any kind.
+    warranties or conditions of any kind. The underlying model is provided by a
+    third party and subject to a separate license, available [here](https://huggingface.co/spaces/bigscience/license).
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -58,6 +63,10 @@ class BloomBackbone(Backbone):
         "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
     }
 
+    # Pretrained BLOOM decoder.
+    model = keras_nlp.models.BloomBackbone.from_preset("bloom_560m_multi")
+    model(input_data)
+
     # Randomly initialized BLOOM decoder with a custom config.
     model = keras_nlp.models.BloomBackbone(
         vocabulary_size=10,
@@ -151,3 +160,7 @@ class BloomBackbone(Backbone):
             }
         )
         return config
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/bloom/bloom_presets.py b/keras_nlp/models/bloom/bloom_presets.py
new file mode 100644
index 00000000..d3e9c780
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_presets.py
@@ -0,0 +1,30 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""BLOOM model preset configurations."""
+
+backbone_presets = {
+    "bloom_560m_multi": {
+        "metadata": {
+            "description": (
+                "24-layer Bloom model. trained on 45 natural languages and "
+                "12 programming languages."
+            ),
+            "params": 816115712,
+            "official_name": "BLOOM",
+            "path": "bloom",
+            "model_card": "https://huggingface.co/bigscience/bloom",
+        },
+        "kaggle_handle": "kaggle://keras/bloom/keras/bloom_560m_multi/1",
+    },
+}
diff --git a/keras_nlp/models/bloom/bloom_tokenizer.py b/keras_nlp/models/bloom/bloom_tokenizer.py
new file mode 100644
index 00000000..cc3fcc2f
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_tokenizer.py
@@ -0,0 +1,123 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.models.bloom.bloom_presets import backbone_presets
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras_nlp_export("keras_nlp.models.BloomTokenizer")
+class BloomTokenizer(BytePairTokenizer):
+    """A BLOOM tokenizer using Byte-Pair Encoding subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.BytePairTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by BLOOM
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a BLOOM preset.
+
+    This tokenizer does not provide truncation or padding of inputs.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space.
+
+    Examples:
+
+    ```python
+    # Unbatched input.
+    tokenizer = keras_nlp.models.BloomTokenizer.from_preset("bloom_560m_multi")
+    tokenizer("The quick brown fox jumped.")
+
+    # Batched input.
+    tokenizer(["The quick brown fox jumped.", "The fox slept."])
+
+    # Detokenization.
+    tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
+
+    # Custom vocabulary.
+    vocab = {"
": 0, "
": 1, "<pad>": 2, "a": 3, "Ġquick": 4, "Ġfox": 5}
+    merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
+    merges += ["Ġ f", "o x", "Ġf ox"]
+    tokenizer = keras_nlp.models.BloomTokenizer(vocabulary=vocab, merges=merges)
+    tokenizer("a quick fox.")
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        **kwargs,
+    ):
+        self.bos_token = "
"
+        self.eos_token = "
"
+        self.pad_token = "<pad>"
+
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            unsplittable_tokens=[
+                self.bos_token,
+                self.eos_token,
+                self.pad_token,
+            ],
+            **kwargs,
+        )
+
+    def set_vocabulary_and_merges(self, vocabulary, merges):
+        super().set_vocabulary_and_merges(vocabulary, merges)
+
+        if vocabulary is not None:
+            # Check for necessary special tokens.
+            for token in [self.bos_token, self.eos_token, self.pad_token]:
+                if token not in self.get_vocabulary():
+                    raise ValueError(
+                        f"Cannot find token `'{token}'` in the provided "
+                        f"`vocabulary`. Please provide `'{token}'` in "
+                        "your `vocabulary` or use a pretrained `vocabulary` name."
+                    )
+
+            self.bos_token_id = self.token_to_id(self.bos_token)
+            self.eos_token_id = self.token_to_id(self.eos_token)
+            self.pad_token_id = self.token_to_id(self.pad_token)
+        else:
+            self.bos_token_id = None
+            self.eos_token_id = None
+            self.pad_token_id = None
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
+
+    def get_config(self):
+        config = super().get_config()
+        # In the constructor, we pass the list of special tokens to the
+        # `unsplittable_tokens` arg of the superclass' constructor. Hence, we
+        # delete it from the config here.
+        del config["unsplittable_tokens"]
+        return config
diff --git a/keras_nlp/models/generative_task.py b/keras_nlp/models/generative_task.py
index 14d4c88f..9a461926 100644
--- a/keras_nlp/models/generative_task.py
+++ b/keras_nlp/models/generative_task.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
+
 import tensorflow as tf
 import tree
 
@@ -80,18 +82,30 @@ class GenerativeTask(Task):
 
             @jax.jit
             def compiled_generate_function(inputs, end_token_id, state):
-                # The only state we update during generation is sampler state,
-                # all weights are fixed and will not change.
-                mapping = zip(self._sampler.variables, state)
+                (
+                    sampler_variables,
+                    trainable_variables,
+                    non_trainable_variables,
+                ) = state
+                mapping = itertools.chain(
+                    zip(self._sampler.variables, sampler_variables),
+                    zip(self.trainable_variables, trainable_variables),
+                    zip(self.non_trainable_variables, non_trainable_variables),
+                )
 
                 with keras.StatelessScope(state_mapping=mapping) as scope:
                     outputs = self.generate_step(inputs, end_token_id)
 
                 # Get updated sampler variables from the stateless scope.
-                state = []
+                sampler_variables = []
                 for v in self._sampler.variables:
                     new_v = scope.get_current_value(v)
-                    state.append(new_v if new_v is not None else v)
+                    sampler_variables.append(new_v if new_v is not None else v)
+                state = (
+                    sampler_variables,
+                    trainable_variables,
+                    non_trainable_variables,
+                )
                 return outputs, state
 
             def wrapped_generate_function(
@@ -99,15 +113,20 @@ class GenerativeTask(Task):
                 end_token_id=None,
             ):
                 # Create an explicit tuple of all variable state.
+                state = (
+                    self._sampler.variables,
+                    self.trainable_variables,
+                    self.non_trainable_variables,
+                )
                 inputs = tree.map_structure(ops.convert_to_tensor, inputs)
                 outputs, state = compiled_generate_function(
                     inputs,
                     end_token_id,
-                    self._sampler.variables,
+                    state,
                 )
                 # Only assign the sampler variables (random seeds), as other
                 # model variables should never be updated in generation.
-                for ref_v, v in zip(self._sampler.variables, state):
+                for ref_v, v in zip(self._sampler.variables, state[0]):
                     ref_v.assign(v)
                 return outputs
 
diff --git a/tools/checkpoint_conversion/convert_bloom_checkpoints.py b/tools/checkpoint_conversion/convert_bloom_checkpoints.py
index 1fc895c5..38acd099 100644
--- a/tools/checkpoint_conversion/convert_bloom_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_bloom_checkpoints.py
@@ -12,51 +12,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
-import torch
-import transformers
-from absl import app
-from absl import flags
-from checkpoint_conversion_utils import get_md5_checksum
+import json
+import os
 
-from keras_nlp.models.bloom.bloom_backbone import BloomBackbone
+os.environ["KERAS_BACKEND"] = "torch"
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+import huggingface_hub  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+import transformers  # noqa: E402
+from absl import app  # noqa: E402
+from absl import flags  # noqa: E402
+
+import keras_nlp  # noqa: E402
+from keras_nlp.models import BloomBackbone  # noqa: E402
+from keras_nlp.models import BloomTokenizer  # noqa: E402
 
 FLAGS = flags.FLAGS
 
 PRESET_MAP = {
-    "bloom_tiny": "bigscience/bloom-560m",
-    "bloom_extra_small": "bigscience/bloom-1b1",
-    "bloom_small": "bigscience/bloom-1b7",
-    "bloom_meduim": "bigscience/bloom-3b",
-    "bloom_large": "bigscience/bloom-7b1",
-    "bloom_extra_large": "bigscience/bloom",
+    "bloom_560m_multi": "bigscience/bloom-560m",
+    "bloom_1.1b_multi": "bigscience/bloom-1b1",
+    "bloom_1.7b_multi": "bigscience/bloom-1b7",
+    "bloom_3b_multi": "bigscience/bloom-3b",
+    "bloom_7b_multi": "bigscience/bloom-7b1",
+    "bloom_176b_multi": "bigscience/bloom",
 }
 
+EXTRACT_DIR = "./model"
+
+
 flags.DEFINE_string(
     "preset", None, f'Must be one of {",".join(PRESET_MAP.keys())}'
 )
 flags.mark_flag_as_required("preset")
 
 
-def convert_checkpoints(hf_model):
+def download_hf_model(hf_model_name):
+    hf_model_dir = huggingface_hub.snapshot_download(
+        repo_id=hf_model_name,
+        allow_patterns=["*.json", "*.bin"],
+        ignore_patterns=["onnx/*"],
+        local_dir=EXTRACT_DIR,
+    )
+
+    return hf_model_dir
+
+
+def convert_model(hf_model):
     # get huggingface model configuration.
     hf_config = hf_model.config.to_dict()
 
-    cfg = {}
-    cfg["vocabulary_size"] = hf_config["vocab_size"]
-    cfg["num_layers"] = hf_config["n_layer"]
-    cfg["num_heads"] = hf_config["n_head"]
-    cfg["hidden_dim"] = hf_config["hidden_size"]
-    cfg["intermediate_dim"] = hf_config["hidden_size"] * 4
-    cfg["dropout"] = hf_config["hidden_dropout"]
-    cfg["layer_norm_epsilon"] = hf_config["layer_norm_epsilon"]
-
-    hidden_dim = cfg["hidden_dim"]
-    num_heads = cfg["num_heads"]
-    head_dim = hidden_dim // num_heads
+    kwargs = {}
+    kwargs["vocabulary_size"] = hf_config["vocab_size"]
+    kwargs["num_layers"] = hf_config["n_layer"]
+    kwargs["num_heads"] = hf_config["n_head"]
+    kwargs["hidden_dim"] = hf_config["hidden_size"]
+    kwargs["intermediate_dim"] = hf_config["hidden_size"] * 4
+    kwargs["dropout"] = hf_config["hidden_dropout"]
+    kwargs["layer_norm_epsilon"] = hf_config["layer_norm_epsilon"]
+
+    return BloomBackbone(**kwargs)
+
+
+def convert_tokenizer(hf_model_dir):
+    tokenizer_file_path = os.path.join(hf_model_dir, "tokenizer.json")
+    with open(tokenizer_file_path) as tokenizer_file:
+        hf_tokenizer = json.load(tokenizer_file)
 
-    # Intialize Bloom model with the weights.
-    keras_model = BloomBackbone(**cfg)
+    vocab = hf_tokenizer["model"]["vocab"]
+    merges = hf_tokenizer["model"]["merges"]
+
+    return BloomTokenizer(vocabulary=vocab, merges=merges)
+
+
+def convert_weights(keras_model, hf_model):
+    hidden_dim = keras_model.hidden_dim
+    num_heads = keras_model.num_heads
+    head_dim = hidden_dim // num_heads
+    num_layers = keras_model.num_layers
 
     # get huggingface model weights.
     hf_wts = hf_model.state_dict()
@@ -78,7 +113,7 @@ def convert_checkpoints(hf_model):
     keras_model.get_layer("final_layernorm").beta.assign(hf_wts["ln_f.bias"])
 
     # Decoder layers.
-    for i in range(cfg["num_layers"]):
+    for i in range(num_layers):
         decoder_layer = keras_model.get_layer(f"transformer_layer_{i}")
         # LayrNorm.
         decoder_layer._pre_attention_layernorm.gamma.assign(
@@ -141,54 +176,72 @@ def convert_checkpoints(hf_model):
             hf_wts[f"h.{i}.mlp.dense_4h_to_h.bias"]
         )
 
-    # Save the model.
-    print(f"\n-> Saving KerasNLP model weights to `{FLAGS.preset}.weights.h5`.")
-    keras_model.save_weights(f"{FLAGS.preset}.weights.h5")
 
-    return keras_model
-
-
-def check_output(keras_model, hf_model):
-    hf_model_input = {
-        "input_ids": torch.tensor([[59414, 15, 2670, 35433, 632, 207595]]),
-        "attention_mask": torch.tensor([[1, 1, 1, 1, 1, 1]]),
-    }
-
-    hf_model_outputs = hf_model(**hf_model_input)
-    hf_model_outputs = hf_model_outputs.last_hidden_state
-    hf_model_outputs = hf_model_outputs.detach().numpy()
+def validate_output(
+    hf_model,
+    keras_model,
+    hf_tokenizer,
+    keras_tokenizer,
+):
+    input_str = ["the quick brown fox ran, galloped and jumped."]
 
+    # KerasNLP
+    token_ids = torch.tensor(keras_tokenizer(input_str))
+    padding_mask = token_ids != 3
     keras_model_input = {
-        "token_ids": torch.tensor([[59414, 15, 2670, 35433, 632, 207595]]),
-        "padding_mask": torch.tensor([[1, 1, 1, 1, 1, 1]]),
+        "token_ids": token_ids,
+        "padding_mask": padding_mask,
     }
-
     keras_model_outputs = keras_model.predict(keras_model_input)
 
-    # Comparing the outputs.
-    print("KerasNLP output:", keras_model_outputs[0, 0, :10])
-    print("HF output:", hf_model_outputs[0, 0, :10])
-    print("Difference:", np.mean(keras_model_outputs - hf_model_outputs))
+    hf_model_input = hf_tokenizer(input_str, return_tensors="pt")
+
+    hf_model_outputs = hf_model(**hf_model_input).last_hidden_state
+    hf_model_outputs = hf_model_outputs.detach().numpy()
 
-    # Show the MD5 checksum of the model weights.
-    print("Model md5sum: ", get_md5_checksum(f"./{FLAGS.preset}.weights.h5"))
+    # Comparing the outputs.
+    print("🔶 KerasNLP output:", keras_model_outputs[0, 0, :10])
+    print("🔶 HF output:", hf_model_outputs[0, 0, :10])
+    print("🔶 Difference:", np.mean(keras_model_outputs - hf_model_outputs))
 
 
 def main(_):
+    preset = FLAGS.preset
+
     assert (
-        FLAGS.preset in PRESET_MAP.keys()
-    ), f'Invalid preset {FLAGS.preset}. Must be one of {",".join(PRESET_MAP.keys())}'
+        preset in PRESET_MAP.keys()
+    ), f'Invalid preset {preset}. Must be one of {",".join(PRESET_MAP.keys())}'
+
+    print(f"✅ Coverting {preset}")
 
-    hf_model_name = PRESET_MAP[FLAGS.preset]
+    hf_model_name = PRESET_MAP[preset]
+    hf_model_dir = download_hf_model(hf_model_name)
+    print("✅ Huggingface model downloaded from hub")
 
-    print("\n-> Loading HF model.")
-    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model = transformers.BloomModel.from_pretrained(hf_model_dir)
+    hf_tokenizer = transformers.BloomTokenizerFast.from_pretrained(hf_model_dir)
+    print("✅ Huggingface model loaded")
 
-    print("\n-> Converting model checkpoint.")
-    keras_model = convert_checkpoints(hf_model)
+    keras_model = convert_model(hf_model)
+    keras_tokenizer = convert_tokenizer(hf_model_dir)
+    print("✅ Keras model loaded")
 
-    print("\n-> Checking keras model output.")
-    check_output(keras_model, hf_model)
+    convert_weights(keras_model, hf_model)
+    print("✅ Weights converted")
+
+    validate_output(
+        hf_model,
+        keras_model,
+        hf_tokenizer,
+        keras_tokenizer,
+    )
+    print("✅ Numerics validated")
+
+    keras_nlp.src.utils.preset_utils.save_to_preset(keras_model, preset)
+    keras_nlp.src.utils.preset_utils.save_to_preset(
+        keras_tokenizer, preset, config_filename="tokenizer.json"
+    )
+    print("✅ Preset saved")
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/keras_nlp/models/bloom/bloom_backbone_test.py b/keras_nlp/models/bloom/bloom_backbone_test.py
index 99cb9d35..83732e49 100644
--- a/keras_nlp/models/bloom/bloom_backbone_test.py
+++ b/keras_nlp/models/bloom/bloom_backbone_test.py
@@ -19,7 +19,7 @@ from keras_nlp.models.bloom.bloom_backbone import BloomBackbone
 from keras_nlp.tests.test_case import TestCase
 
 
-class BloomTest(TestCase):
+class BloomBackboneTest(TestCase):
     def setUp(self):
         self.init_kwargs = {
             "vocabulary_size": 10,
@@ -49,3 +49,28 @@ class BloomTest(TestCase):
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
         )
+
+    @pytest.mark.large
+    def test_smallest_preset(self):
+        self.run_preset_test(
+            cls=BloomBackbone,
+            preset="bloom_560m_multi",
+            input_data={
+                "token_ids": ops.array([[101, 1996, 4248, 102]], dtype="int32"),
+                "padding_mask": ops.ones((1, 4), dtype="int32"),
+            },
+            expected_output_shape=(1, 4, 1024),
+            # The forward pass from a preset should be stable!
+            expected_partial_output=ops.array(
+                [2.4394186, 1.4131186, -2.7810357, -6.330823, -1.0599766]
+            ),
+        )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in BloomBackbone.presets:
+            self.run_preset_test(
+                cls=BloomBackbone,
+                preset=preset,
+                input_data=self.input_data,
+            )
diff --git a/keras_nlp/models/bloom/bloom_tokenizer_test.py b/keras_nlp/models/bloom/bloom_tokenizer_test.py
new file mode 100644
index 00000000..9ae9c0cc
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_tokenizer_test.py
@@ -0,0 +1,63 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from keras_nlp.models.bloom.bloom_tokenizer import BloomTokenizer
+from keras_nlp.tests.test_case import TestCase
+
+
+class BloomTokenizerTest(TestCase):
+    def setUp(self):
+        self.vocab = ["!", "air", "Ġair", "plane", "Ġat", "port"]
+        self.vocab += ["
", "
", "<pad>"]
+        self.vocab = dict([(token, i) for i, token in enumerate(self.vocab)])
+        self.merges = ["Ġ a", "Ġ t", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        self.merges += ["Ġa t", "p o", "r t", "Ġt h", "ai r", "pl a", "po rt"]
+        self.merges += ["Ġai r", "Ġa i", "pla ne"]
+        self.init_kwargs = {"vocabulary": self.vocab, "merges": self.merges}
+        self.input_data = [
+            "
airplane at airport<pad>",
+            "
 airplane airport<pad>",
+        ]
+
+    def test_tokenizer_basics(self):
+        self.run_preprocessing_layer_test(
+            cls=BloomTokenizer,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output=[[6, 1, 3, 4, 2, 5, 8], [6, 2, 3, 2, 5, 8]],
+        )
+
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            BloomTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
+    @pytest.mark.large
+    def test_smallest_preset(self):
+        self.run_preset_test(
+            cls=BloomTokenizer,
+            preset="bloom_560m_multi",
+            input_data=["The quick brown fox."],
+            expected_output=[[2175, 23714, 73173, 144252, 17]],
+        )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in BloomTokenizer.presets:
+            self.run_preset_test(
+                cls=BloomTokenizer,
+                preset=preset,
+                input_data=self.input_data,
+            )