Repository: keras-team/keras-nlp

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.


Discussion:
Thanks for the report! Will add the new presets.

Resolving patch:
diff --git a/keras_nlp/models/mistral/mistral_presets.py b/keras_nlp/models/mistral/mistral_presets.py
--- a/keras_nlp/models/mistral/mistral_presets.py
+++ b/keras_nlp/models/mistral/mistral_presets.py
@@ -35,4 +35,14 @@
         },
         "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
--- a/keras_nlp/models/mistral/mistral_transformer_decoder.py
+++ b/keras_nlp/models/mistral/mistral_transformer_decoder.py
@@ -207,17 +207,20 @@ def _compute_self_attention_mask(
             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/tools/checkpoint_conversion/convert_mistral_checkpoints.py b/tools/checkpoint_conversion/convert_mistral_checkpoints.py
--- a/tools/checkpoint_conversion/convert_mistral_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_mistral_checkpoints.py
@@ -33,6 +33,7 @@
 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: