Repository: keras-team/keras-nlp (python)

Issue:
Add Mistral 0.2 models as possible presets
**Is your feature request related to a problem? Please describe.**
<!---
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
-->
We can currently load Mistral 7b models with `keras_nlp.models.MistralCausalLM.from_preset("mistral_7b_en")` (or `mistral_instruct_7b_en`).
I noticed those are the version `0.1` of the models. The versions `0.2` have significantly improved while using the same code base.

**Describe the solution you'd like**
<!---
A clear and concise description of what you want to happen.
-->
It would be nice to offer the `0.2` variants of both base and instruct models as possible presets.

**Describe alternatives you've considered**
<!---
A clear and concise description of any alternative solutions or features you've considered.
-->
Just using the `0.1` versions or other libraries… However Keras offers nice integration with JAX + sharding!

**Additional context**
<!---
Add any other context or screenshots about the feature request here.
-->
I want to use Keras models as backbones to train VLM’s and Mistral is a very strong variant at 7B.

Pull request description:
Add Mistral Instruct V0.2 preset
Closes #1515 

Adds Mistral Instruct v0.2 as a preset. The model has been uploaded on Kaggle.

~Just need to test the generator to make sure HF and Keras outputs match.~

Patch:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 30f8a53b..407a4b7a 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -26,5 +26,6 @@ from keras_nlp import models
 from keras_nlp import samplers
 from keras_nlp import tokenizers
 from keras_nlp import utils
+from keras_nlp.utils.preset_utils import upload_preset
 from keras_nlp.version_utils import __version__
 from keras_nlp.version_utils import version
diff --git a/keras_nlp/models/backbone.py b/keras_nlp/models/backbone.py
index bfdc8207..08b9f86e 100644
--- a/keras_nlp/models/backbone.py
+++ b/keras_nlp/models/backbone.py
@@ -17,6 +17,7 @@ from keras_nlp.backend import config
 from keras_nlp.backend import keras
 from keras_nlp.utils.preset_utils import check_preset_class
 from keras_nlp.utils.preset_utils import load_from_preset
+from keras_nlp.utils.preset_utils import save_to_preset
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
@@ -141,6 +142,14 @@ class Backbone(keras.Model):
             config_overrides=kwargs,
         )
 
+    def save_to_preset(self, preset):
+        """Save backbone to a preset directory.
+
+        Args:
+            preset: The path to the local model preset directory.
+        """
+        save_to_preset(self, preset)
+
     def __init_subclass__(cls, **kwargs):
         # Use __init_subclass__ to setup a correct docstring for from_preset.
         super().__init_subclass__(**kwargs)
diff --git a/keras_nlp/models/mistral/mistral_causal_lm.py b/keras_nlp/models/mistral/mistral_causal_lm.py
index e9bd4e56..800df1c8 100644
--- a/keras_nlp/models/mistral/mistral_causal_lm.py
+++ b/keras_nlp/models/mistral/mistral_causal_lm.py
@@ -215,6 +215,132 @@ class MistralCausalLM(GenerativeTask):
             "padding_mask": padding_mask,
         }
 
+    def score(
+        self,
+        token_ids,
+        padding_mask=None,
+        scoring_mode="logits",
+        layer_intercept_fn=None,
+        target_ids=None,
+    ):
+        """Score a generation represented by the provided token ids.
+
+        Args:
+            token_ids: A <int>[batch_size, num_tokens] tensor containing tokens
+                to score. Typically, this tensor captures the output from a call
+                to `MistralCausalLM.generate()`, i.e., tokens for both the input
+                text and the model-generated text.
+            padding_mask: A <bool>[batch_size, num_tokens] tensor indicating the
+                tokens that should be preserved during generation. This is an
+                artifact required by the MistralBackbone and isn't influential
+                on the computation of this function. If omitted, this function
+                uses `keras.ops.ones()` to create a tensor of the appropriate
+                shape.
+            scoring_mode: The type of scores to return, either "logits" or
+                "loss", both will be per input token.
+            layer_intercept_fn: An optional function for augmenting activations
+                with additional computation, for example, as part of
+                interpretability research. This function will be passed the
+                activations as its first parameter and a numeric index
+                associated with that backbone layer. _This index _is not_ an
+                index into `self.backbone.layers`. The index -1 accompanies the
+                embeddings returned by calling `self.backbone.token_embedding()`
+                on `token_ids` in the forward direction. All subsequent indexes
+                will be 0-based indices for the activations returned by each of
+                the Transformers layers in the backbone. This function must
+                return a <float>[batch_size, num_tokens, hidden_dims] tensor
+                that can be passed as an input to the next layer in the model.
+            target_ids: An <bool>[batch_size, num_tokens] tensor containing the
+                predicted tokens against which the loss should be computed. If a
+                span of tokens is provided (sequential truthy values along
+                axis=1 in the tensor), the loss will be computed as the
+                aggregate across those tokens.
+
+        Raises:
+            ValueError: If an unsupported scoring_mode is provided, or if the
+                target_ids are not provided when using ScoringMode.LOSS.
+
+        Returns:
+            The per-token scores as a tensor of size
+            <float>[batch_size, num_tokens, vocab_size] in "logits" mode, or
+            <float>[batch_size, num_tokens] in "loss" mode.
+
+        Examples:
+
+        Compute gradients between embeddings and loss scores with TensorFlow:
+        ```python
+        mistral_lm = keras_nlp.models.MistralCausalLM.from_preset(
+            "mistral_7b_en"
+        )
+        generations = mistral_lm.generate(
+            ["This is a", "Where are you"],
+            max_length=30
+        )
+        preprocessed = mistral_lm.preprocessor.generate_preprocess(generations)
+        generation_ids = preprocessed["token_ids"]
+        padding_mask = preprocessed["padding_mask"]
+        target_ids = keras.ops.roll(generation_ids, shift=-1, axis=1)
+
+        embeddings = None
+        with tf.GradientTape(watch_accessed_variables=True) as tape:
+            def layer_intercept_fn(x, i):
+                if i == -1:
+                    nonlocal embeddings, tape
+                    embeddings = x
+                    tape.watch(embeddings)
+                return x
+
+            losses = mistral_lm.score(
+                token_ids=generation_ids,
+                padding_mask=padding_mask,
+                scoring_mode="loss",
+                layer_intercept_fn=layer_intercept_fn,
+                target_ids=target_ids,
+            )
+
+        grads = tape.gradient(losses, embeddings)
+        ```
+        """
+        if scoring_mode not in ("logits", "loss"):
+            raise ValueError(
+                "Unsupported scoring_mode. Must be one of 'logits' or 'loss'."
+            )
+
+        if scoring_mode == "loss" and target_ids is None:
+            raise ValueError(
+                "Cannot compute loss without targets. Please provide target "
+                "token ids via the target_ids parameter."
+            )
+
+        batch_shape = ops.shape(token_ids)[:2]
+        assert len(batch_shape) == 2
+
+        if layer_intercept_fn is None:
+
+            def default_layer_intercept_fn(x, unused_i):
+                return x
+
+            layer_intercept_fn = default_layer_intercept_fn
+
+        token_embeddings = self.backbone.token_embedding(token_ids)
+        x = layer_intercept_fn(token_embeddings, -1)
+
+        for i, transformer_layer in enumerate(self.backbone.transformer_layers):
+            x = transformer_layer(x, decoder_padding_mask=padding_mask)
+            x = layer_intercept_fn(x, i)
+
+        x = self.backbone.layer_norm(x)
+        logits = self.backbone.token_embedding(x, reverse=True)
+
+        if scoring_mode == "logits":
+            return logits
+
+        per_token_loss_fn = keras.losses.SparseCategoricalCrossentropy(
+            from_logits=True, reduction="none"
+        )
+        per_token_loss = per_token_loss_fn(target_ids, logits)
+        return per_token_loss
+
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/mistral/mistral_presets.py b/keras_nlp/models/mistral/mistral_presets.py
index 7fb4b4e0..fdee3963 100644
--- a/keras_nlp/models/mistral/mistral_presets.py
+++ b/keras_nlp/models/mistral/mistral_presets.py
@@ -35,4 +35,14 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/6",
     },
+    "mistral_0.2_instruct_7b_en": {
+        "metadata": {
+            "description": "Mistral 7B instruct Version 0.2 model",
+            "params": 7241732096,
+            "official_name": "Mistral",
+            "path": "mistral",
+            "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
+        },
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.2_instruct_7b_en/1",
+    },
 }
diff --git a/keras_nlp/models/mistral/mistral_transformer_decoder.py b/keras_nlp/models/mistral/mistral_transformer_decoder.py
index 7c90ab91..36b7f594 100644
--- a/keras_nlp/models/mistral/mistral_transformer_decoder.py
+++ b/keras_nlp/models/mistral/mistral_transformer_decoder.py
@@ -207,17 +207,20 @@ class MistralTransformerDecoder(keras.layers.Layer):
             else self_attention_cache_update_index
         )
 
-        # Mistral uses a banded attention mask
-        causal_mask_lower = compute_causal_mask(
+        # The lower traingular attention mask
+        causal_mask = compute_causal_mask(
             batch_size, input_length, output_length, cache_update_index
         )
-        # Below is a workaround for `ops.triu` for Keras 2.
-        # TODO(tirthasheshpatel): Use `ops.triu` once Keras 2 support is removed.
-        # causal_mask = ops.triu(causal_mask_lower, k=-self.sliding_window)
-        i = ops.arange(output_length)[:, None] + cache_update_index
-        j = ops.arange(input_length)[None, :]
-        causal_mask_upper = ops.cast(i < j + self.sliding_window, "int32")
-        causal_mask = ops.minimum(causal_mask_lower, causal_mask_upper)
+
+        # Mistral uses a banded attention mask if sliding window is not None
+        if self.sliding_window is not None:
+            # Below is a workaround for `ops.triu` for Keras 2.
+            # TODO(tirthasheshpatel): Use `ops.triu` once Keras 2 support is removed.
+            # causal_mask = ops.triu(causal_mask, k=-self.sliding_window)
+            i = ops.arange(output_length)[:, None] + cache_update_index
+            j = ops.arange(input_length)[None, :]
+            causal_mask_upper = ops.cast(i < j + self.sliding_window, "int32")
+            causal_mask = ops.minimum(causal_mask, causal_mask_upper)
 
         return (
             ops.minimum(decoder_mask, causal_mask)
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
index 4c26e452..834b99e5 100644
--- a/keras_nlp/tokenizers/tokenizer.py
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -18,6 +18,8 @@ from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.layers.preprocessing.preprocessing_layer import (
     PreprocessingLayer,
 )
+from keras_nlp.utils.preset_utils import TOKENIZER_CONFIG_FILE
+from keras_nlp.utils.preset_utils import save_to_preset
 
 
 @keras_nlp_export("keras_nlp.tokenizers.Tokenizer")
@@ -121,5 +123,13 @@ class Tokenizer(PreprocessingLayer):
             f"{self.__class__.__name__}."
         )
 
+    def save_to_preset(self, preset):
+        """Save tokenizer to a preset directory.
+
+        Args:
+            preset: The path to the local model preset directory.
+        """
+        save_to_preset(self, preset, config_filename=TOKENIZER_CONFIG_FILE)
+
     def call(self, inputs, *args, training=None, **kwargs):
         return self.tokenize(inputs, *args, **kwargs)
diff --git a/keras_nlp/utils/preset_utils.py b/keras_nlp/utils/preset_utils.py
index 01c11a3d..dcee9bc6 100644
--- a/keras_nlp/utils/preset_utils.py
+++ b/keras_nlp/utils/preset_utils.py
@@ -16,6 +16,9 @@ import datetime
 import json
 import os
 
+from absl import logging
+
+from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import config as backend_config
 from keras_nlp.backend import keras
 
@@ -27,6 +30,8 @@ except ImportError:
 KAGGLE_PREFIX = "kaggle://"
 GS_PREFIX = "gs://"
 TOKENIZER_ASSET_DIR = "assets/tokenizer"
+CONFIG_FILE = "config.json"
+TOKENIZER_CONFIG_FILE = "tokenizer.json"
 
 
 def get_file(preset, path):
@@ -155,6 +160,115 @@ def save_to_preset(
             metadata_file.write(json.dumps(metadata, indent=4))
 
 
+def _validate_tokenizer(preset, allow_incomplete=False):
+    config_path = get_file(preset, TOKENIZER_CONFIG_FILE)
+    if not os.path.exists(config_path):
+        if allow_incomplete:
+            logging.warning(
+                f"`{TOKENIZER_CONFIG_FILE}` is missing from the preset directory `{preset}`."
+            )
+            return
+        else:
+            raise FileNotFoundError(
+                f"`{TOKENIZER_CONFIG_FILE}` is missing from the preset directory `{preset}`. "
+                "To upload the model without a tokenizer, "
+                "set `allow_incomplete=True`."
+            )
+    try:
+        with open(config_path) as config_file:
+            config = json.load(config_file)
+    except Exception as e:
+        raise ValueError(
+            f"Tokenizer config file `{config_path}` is an invalid json file. "
+            f"Error message: {e}"
+        )
+    layer = keras.saving.deserialize_keras_object(config)
+
+    if not config["assets"]:
+        raise ValueError(
+            f"Tokenizer config file {config_path} is missing `asset`."
+        )
+
+    for asset in config["assets"]:
+        asset_path = os.path.join(preset, asset)
+        if not os.path.exists(asset_path):
+            raise FileNotFoundError(
+                f"Asset `{asset}` doesn't exist in the preset direcotry `{preset}`."
+            )
+    config_dir = os.path.dirname(config_path)
+    asset_dir = os.path.join(config_dir, TOKENIZER_ASSET_DIR)
+
+    tokenizer = get_tokenizer(layer)
+    if not tokenizer:
+        raise ValueError(f"Model or layer `{layer}` is missing tokenizer.")
+    tokenizer.load_assets(asset_dir)
+
+
+def _validate_backbone(preset):
+    config_path = os.path.join(preset, CONFIG_FILE)
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(
+            f"`{CONFIG_FILE}` is missing from the preset directory `{preset}`."
+        )
+    try:
+        with open(config_path) as config_file:
+            config = json.load(config_file)
+    except Exception as e:
+        raise ValueError(
+            f"Config file `{config_path}` is an invalid json file. "
+            f"Error message: {e}"
+        )
+
+    if config["weights"]:
+        weights_path = os.path.join(preset, config["weights"])
+        if not os.path.exists(weights_path):
+            raise FileNotFoundError(
+                f"The weights file doesn't exist in preset directory `{preset}`."
+            )
+    else:
+        raise ValueError(
+            f"No weights listed in `{CONFIG_FILE}`. Make sure to use "
+            "`save_to_preset()` which adds additional data to a serialized "
+            "Keras object."
+        )
+
+
+@keras_nlp_export("keras_nlp.upload_preset")
+def upload_preset(
+    uri,
+    preset,
+    allow_incomplete=False,
+):
+    """Upload a preset directory to a model hub.
+
+    Args:
+        uri: The URI identifying model to upload to.
+             URIs with format
+             `kaggle://<KAGGLE_USERNAME>/<MODEL>/<FRAMEWORK>/<VARIATION>`
+             will be uploaded to Kaggle Hub.
+        preset: The path to the local model preset directory.
+        allow_incomplete: If True, allows the upload of presets without
+                          a tokenizer configuration. Otherwise, a tokenizer
+                          is required.
+    """
+
+    # Check if preset directory exists.
+    if not os.path.exists(preset):
+        raise FileNotFoundError(f"The preset directory {preset} doesn't exist.")
+
+    _validate_backbone(preset)
+    _validate_tokenizer(preset, allow_incomplete)
+
+    if uri.startswith(KAGGLE_PREFIX):
+        kaggle_handle = uri.removeprefix(KAGGLE_PREFIX)
+        kagglehub.model_upload(kaggle_handle, preset)
+    else:
+        raise ValueError(
+            f"Unexpected URI `'{uri}'`. Kaggle upload format should follow "
+            "`kaggle://<KAGGLE_USERNAME>/<MODEL>/<FRAMEWORK>/<VARIATION>`."
+        )
+
+
 def load_from_preset(
     preset,
     load_weights=True,
diff --git a/tools/checkpoint_conversion/convert_mistral_checkpoints.py b/tools/checkpoint_conversion/convert_mistral_checkpoints.py
index 7e13b9dd..ae3b1f5b 100644
--- a/tools/checkpoint_conversion/convert_mistral_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_mistral_checkpoints.py
@@ -33,6 +33,7 @@ from keras_nlp.utils.preset_utils import save_to_preset
 PRESET_MAP = {
     "mistral_7b_en": "mistralai/Mistral-7B-v0.1",
     "mistral_instruct_7b_en": "mistralai/Mistral-7B-Instruct-v0.1",
+    "mistral_0.2_instruct_7b_en": "mistralai/Mistral-7B-Instruct-v0.2",
 }
 
 FLAGS = flags.FLAGS


Test patch:
diff --git a/keras_nlp/models/mistral/mistral_causal_lm_test.py b/keras_nlp/models/mistral/mistral_causal_lm_test.py
index 3f9d7fab..13f0dad9 100644
--- a/keras_nlp/models/mistral/mistral_causal_lm_test.py
+++ b/keras_nlp/models/mistral/mistral_causal_lm_test.py
@@ -128,3 +128,88 @@ class MistralCausalLMTest(TestCase):
                 preset=preset,
                 input_data=self.input_data,
             )
+
+    def test_score_logits(self):
+        # Setup prompts, models, and associated expected shapes.
+        prompts = ["the quick brown fox", "the quick brown fox"]
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        expected_score_shape = (2, 8, 10)
+
+        # Preprocess prompts to get tokenized representations and padding masks.
+        preprocessed_prompts = causal_lm.preprocessor.generate_preprocess(
+            prompts
+        )
+        token_ids = preprocessed_prompts["token_ids"]
+        padding_mask = preprocessed_prompts["padding_mask"]
+
+        # Get the scores and assert their shape.
+        scores = causal_lm.score(
+            token_ids=token_ids,
+            padding_mask=padding_mask,
+            scoring_mode="logits",
+        )
+
+        self.assertEqual(ops.shape(scores), expected_score_shape)
+
+    def test_score_loss(self):
+        # Setup prompts, models, and associated expected shapes.
+        prompts = ["the quick brown fox", "the quick brown fox"]
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        expected_score_shape = (2, 8)
+
+        # Preprocess prompts to get tokenized representations and padding masks.
+        preprocessed_prompts = causal_lm.preprocessor.generate_preprocess(
+            prompts
+        )
+        token_ids = preprocessed_prompts["token_ids"]
+        padding_mask = preprocessed_prompts["padding_mask"]
+        target_ids = ops.roll(token_ids, shift=-1, axis=1)
+
+        # Get the scores and assert their shape.
+        scores = causal_lm.score(
+            token_ids=token_ids,
+            padding_mask=padding_mask,
+            scoring_mode="loss",
+            target_ids=target_ids,
+        )
+
+        self.assertEqual(ops.shape(scores), expected_score_shape)
+
+    def test_score_layer_intercept_fn_exfiltration(self):
+        # Setup prompts, models, and associated expected shapes.
+        prompts = ["the quick brown fox", "the quick brown fox"]
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        expected_embedded_shape = (2, 8, 8)
+        expected_score_shape = (2, 8, 10)
+
+        # Preprocess prompts to get tokenized representations and padding masks.
+        preprocessed_prompts = causal_lm.preprocessor.generate_preprocess(
+            prompts
+        )
+        token_ids = preprocessed_prompts["token_ids"]
+        padding_mask = preprocessed_prompts["padding_mask"]
+
+        # Setup a custom intercept function that extracts the embeddings to a
+        # a variable from the embeddings layer and otherwise asserts on shapes.
+        embedded_prompts = None
+
+        def layer_intercept_fn_for_testing(x, i):
+            if i == -1:
+                nonlocal embedded_prompts
+                embedded_prompts = x
+            else:
+                nonlocal expected_embedded_shape
+                self.assertEqual(ops.shape(x), expected_embedded_shape)
+            return x
+
+        # Get the scores.
+        scores = causal_lm.score(
+            token_ids=token_ids,
+            padding_mask=padding_mask,
+            scoring_mode="logits",
+            layer_intercept_fn=layer_intercept_fn_for_testing,
+        )
+
+        # Assert shapes for info exfiltrated into the parent context.
+        self.assertEqual(ops.shape(embedded_prompts), expected_embedded_shape)
+        self.assertEqual(ops.shape(scores), expected_score_shape)
diff --git a/keras_nlp/utils/preset_utils_test.py b/keras_nlp/utils/preset_utils_test.py
index 44dc39f4..289e13b6 100644
--- a/keras_nlp/utils/preset_utils_test.py
+++ b/keras_nlp/utils/preset_utils_test.py
@@ -18,6 +18,7 @@ import os
 import pytest
 from absl.testing import parameterized
 
+from keras_nlp import upload_preset
 from keras_nlp.models.albert.albert_classifier import AlbertClassifier
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.bert.bert_classifier import BertClassifier
@@ -105,3 +106,14 @@ class PresetUtilsTest(TestCase):
 
         with self.assertRaisesRegex(ValueError, "Unknown preset identifier"):
             AlbertClassifier.from_preset("snaggle://bort/bort/bort")
+
+    def test_upload_empty_preset(self):
+        temp_dir = self.get_temp_dir()
+        empty_preset = os.path.join(temp_dir, "empty")
+        os.mkdir(empty_preset)
+        uri = "kaggle://test/test/test"
+
+        with self.assertRaises(FileNotFoundError):
+            upload_preset(uri, empty_preset)
+
+    # TODO: add more test to cover various invalid scenarios such as invalid json, missing files, etc.