Add Microsoft Harrier Model to Hub (keras-team#2764)

* Add microsoft/harrier-oss-v1-0.6b Model to Hub

* Fix gemini code assist review comments

* Fix gemini code assist review comments and colab error

* Fix gemini code assist review comments

* Fix gemini code assist review comments

* Add qwen3 text embedder and preprossor

* Fix gemini code assist review comments

* Update dtype to float32 for numerics match

* Remove gemini assist safe check code suggestions

* Add harrier embedding oss 270m model support

* Fix Gemini Code assist review comments

* Fix review comments for gemma3 embedding models

* Fix Gemini Code assist review comments

* Update Harrier Kaggle Handles

* Fix preset name for harrier model

---------

Co-authored-by: Vijay Dabur <vijay_dabur@epam.com>

Author: vjdaburdev

keras_hub/api/models/__init__.py

@@ -322,6 +322,12 @@
 from keras_hub.src.models.gemma3.gemma3_causal_lm_preprocessor import (
     Gemma3CausalLMPreprocessor as Gemma3CausalLMPreprocessor,
 )
+from keras_hub.src.models.gemma3.gemma3_text_embedder import (
+    Gemma3TextEmbedder as Gemma3TextEmbedder,
+)
+from keras_hub.src.models.gemma3.gemma3_text_embedder_preprocessor import (
+    Gemma3TextEmbedderPreprocessor as Gemma3TextEmbedderPreprocessor,
+)
 from keras_hub.src.models.gemma3.gemma3_tokenizer import (
     Gemma3Tokenizer as Gemma3Tokenizer,
 )
@@ -609,6 +615,12 @@
 from keras_hub.src.models.qwen3.qwen3_causal_lm_preprocessor import (
     Qwen3CausalLMPreprocessor as Qwen3CausalLMPreprocessor,
 )
+from keras_hub.src.models.qwen3.qwen3_text_embedder import (
+    Qwen3TextEmbedder as Qwen3TextEmbedder,
+)
+from keras_hub.src.models.qwen3.qwen3_text_embedder_preprocessor import (
+    Qwen3TextEmbedderPreprocessor as Qwen3TextEmbedderPreprocessor,
+)
 from keras_hub.src.models.qwen3.qwen3_tokenizer import (
     Qwen3Tokenizer as Qwen3Tokenizer,
 )

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -303,4 +303,18 @@
         },
         "kaggle_handle": "kaggle://keras/embeddinggemma/keras/embedding_gemma3_300m/2",
     },
+    # Microsoft harrier model presets
+    "harrier_embedding_oss_270m": {
+        "metadata": {
+            "description": (
+                "Microsoft harrier-oss-v1 270M multilingual text embedding "
+                "model based on the Gemma3-270M architecture, fine-tuned for "
+                "dense retrieval and semantic similarity across 94+ languages. "
+                "Achieves 66.5 on Multilingual MTEB v2."
+            ),
+            "params": 268098176,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/harrier/keras/harrier_embedding_oss_270m/1",
+    },
 }

keras_hub/src/models/gemma3/gemma3_text_embedder.py

@@ -0,0 +1,200 @@
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.gemma3.gemma3_backbone import Gemma3Backbone
+from keras_hub.src.models.gemma3.gemma3_text_embedder_preprocessor import (
+    Gemma3TextEmbedderPreprocessor,
+)
+from keras_hub.src.models.text_embedder import TextEmbedder
+
+
+@keras_hub_export("keras_hub.models.Gemma3TextEmbedder")
+class Gemma3TextEmbedder(TextEmbedder):
+    """An end-to-end Gemma3 model for generating sentence embeddings.
+
+    This model attaches a pooling and optional L2 normalization head to a
+    `keras_hub.models.Gemma3Backbone` instance, mapping from the backbone
+    outputs to fixed-size sentence embeddings suitable for semantic
+    similarity, clustering, and retrieval tasks.
+
+    This is the architecture used by Microsoft's harrier-oss-v1-270m embedding
+    model. For usage of this model with pre-trained weights, use the
+    `from_preset()` constructor.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to raw inputs during
+    `fit()`, `predict()`, and `evaluate()`. This is done by default when
+    creating the model with `from_preset()`.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:
+        backbone: A `keras_hub.models.Gemma3Backbone` instance.
+        preprocessor: A `keras_hub.models.Gemma3TextEmbedderPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
+        pooling_mode: string. The pooling strategy to use. One of `"last"` or
+            `"mean"`. Defaults to `"last"`.
+            - `"last"`: Use the hidden state of the last non-padding token.
+              This is the standard pooling strategy for decoder-only models
+              and matches the approach used by harrier-oss-v1-270m.
+            - `"mean"`: Attention-mask-aware mean pooling over all tokens.
+        normalize: bool. Whether to L2 normalize the output embeddings.
+            Defaults to `True`.
+
+    Examples:
+
+    Raw string data.
+    ```python
+    embedder = keras_hub.models.Gemma3TextEmbedder.from_preset(
+        "hf://microsoft/harrier-oss-v1-270m",
+    )
+
+    # Semantic search.
+    query = "Which planet is known as the Red Planet?"
+    documents = [
+        "Mars is often referred to as the Red Planet.",
+        "Venus is often called Earth's twin.",
+    ]
+    q_emb = embedder.encode_text(
+        "Instruct: Given a web search query, retrieve relevant passages\n"
+        "Query: " + query
+    )
+    d_embs = embedder.encode_text(documents)
+    sims = embedder.similarity(q_emb, d_embs)
+    print("Best match:", documents[sims.numpy().argmax()])
+    ```
+
+    Preprocessed integer data.
+    ```python
+    features = {
+        "token_ids": np.ones(shape=(2, 12), dtype="int32"),
+        "padding_mask": np.ones(shape=(2, 12), dtype="int32"),
+    }
+
+    embedder = keras_hub.models.Gemma3TextEmbedder.from_preset(
+        "hf://microsoft/harrier-oss-v1-270m",
+        preprocessor=None,
+    )
+    embeddings = embedder.predict(features)
+    ```
+
+    Reference:
+     - [harrier-oss-v1 model page](https://huggingface.co/microsoft/harrier-oss-v1-270m)
+    """
+
+    backbone_cls = Gemma3Backbone
+    preprocessor_cls = Gemma3TextEmbedderPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        preprocessor=None,
+        pooling_mode="last",
+        normalize=True,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
+        inputs = backbone.input
+        backbone_output = backbone(inputs)
+        padding_mask = inputs["padding_mask"]
+
+        if (
+            isinstance(backbone_output, dict)
+            and "pooled_output" in backbone_output
+        ):
+            pooled = backbone_output["pooled_output"]
+        else:
+            if isinstance(backbone_output, dict):
+                sequence_output = backbone_output["sequence_output"]
+            else:
+                sequence_output = backbone_output
+
+            # Apply pooling.
+            if pooling_mode == "last":
+                pooled = self._last_token_pooling(sequence_output, padding_mask)
+            elif pooling_mode == "mean":
+                pooled = self._mean_pooling(sequence_output, padding_mask)
+            else:
+                raise ValueError(
+                    f"Invalid pooling_mode: '{pooling_mode}'. "
+                    "Expected one of 'last' or 'mean'."
+                )
+
+            # Apply L2 normalization.
+            if normalize:
+                pooled = self._l2_normalize(pooled)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=pooled,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.pooling_mode = pooling_mode
+        self.normalize = normalize
+
+    @staticmethod
+    def _last_token_pooling(sequence_output, padding_mask):
+        """Pool the hidden state of the last non-padding token."""
+        mask = ops.cast(padding_mask, sequence_output.dtype)
+        # Shift mask left to find the last real token (where mask=1, next=0).
+        mask_shifted = ops.pad(mask, [[0, 0], [0, 1]])[:, 1:]
+        last_token_mask = mask * (1.0 - mask_shifted)
+        return ops.sum(
+            sequence_output * ops.expand_dims(last_token_mask, axis=-1),
+            axis=1,
+        )
+
+    @staticmethod
+    def _mean_pooling(sequence_output, padding_mask):
+        """Attention-mask-aware mean pooling over non-padding tokens."""
+        mask = ops.cast(
+            ops.expand_dims(padding_mask, axis=-1), sequence_output.dtype
+        )
+        sum_embeddings = ops.sum(sequence_output * mask, axis=1)
+        sum_mask = ops.maximum(ops.sum(mask, axis=1), 1e-5)
+        return sum_embeddings / sum_mask
+
+    @staticmethod
+    def _l2_normalize(embeddings):
+        """L2-normalize embeddings along the last axis."""
+        return ops.nn.normalize(embeddings, axis=-1, order=2)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "pooling_mode": self.pooling_mode,
+                "normalize": self.normalize,
+            }
+        )
+        return config
+
+    def similarity(self, query_embeddings, document_embeddings):
+        """Compute similarity between query and document embeddings.
+
+        Computes the dot product between query and document embeddings.
+        When embeddings are L2 normalized (the default), this is
+        equivalent to cosine similarity. Returns a similarity matrix
+        of shape `(num_queries, num_documents)`.
+
+        Args:
+            query_embeddings: An array or tensor of shape
+                `(num_queries, embedding_dim)`.
+            document_embeddings: An array or tensor of shape
+                `(num_documents, embedding_dim)`.
+
+        Returns:
+            A tensor of shape `(num_queries, num_documents)` containing
+            similarity scores.
+        """
+        query_tensor = ops.convert_to_tensor(query_embeddings)
+        document_tensor = ops.convert_to_tensor(document_embeddings)
+        return ops.matmul(query_tensor, ops.transpose(document_tensor))

keras_hub/src/models/gemma3/gemma3_text_embedder_preprocessor.py

@@ -0,0 +1,88 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.gemma3.gemma3_backbone import Gemma3Backbone
+from keras_hub.src.models.gemma3.gemma3_tokenizer import Gemma3Tokenizer
+from keras_hub.src.models.text_embedder_preprocessor import (
+    TextEmbedderPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.Gemma3TextEmbedderPreprocessor")
+class Gemma3TextEmbedderPreprocessor(TextEmbedderPreprocessor):
+    """A Gemma3 preprocessing layer which tokenizes and packs inputs for
+    sentence embedding.
+
+    This preprocessing layer will do three things:
+
+    1. Tokenize any number of input segments using the `tokenizer`.
+    2. Pack the inputs together using a `keras_hub.layers.MultiSegmentPacker`
+       with the appropriate `<bos>`, `<eos>`, and `<pad>` tokens.
+    3. Construct a dictionary with keys `"token_ids"` and `"padding_mask"`
+       that can be passed directly to a `Gemma3TextEmbedder` model.
+
+    This layer can be used directly with `tf.data.Dataset.map` to preprocess
+    string data in the `(x, y, sample_weight)` format used by
+    `keras.Model.fit`.
+
+    Args:
+        tokenizer: A `keras_hub.models.Gemma3Tokenizer` instance.
+        sequence_length: int. The length of the packed inputs. Defaults to
+            `256`.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
+            - `"round_robin"`: Available space is assigned one token at a
+                time in a round-robin fashion to the inputs that still need
+                some, until the limit is reached.
+            - `"waterfall"`: The allocation of the budget is done using a
+                "waterfall" algorithm that allocates quota in a
+                left-to-right manner and fills up the buckets until we run
+                out of budget. It supports an arbitrary number of segments.
+
+    Call arguments:
+        x: A tensor of single string sequences, or a tuple of multiple
+            tensor sequences to be packed together. Inputs may be batched or
+            unbatched. For single sequences, raw python inputs will be
+            converted to tensors. For multiple sequences, pass tensors
+            directly.
+        y: Any label data. Will be passed through unaltered.
+        sample_weight: Any label weight data. Will be passed through
+            unaltered.
+
+    Examples:
+
+    Directly calling the layer on data.
+    ```python
+    preprocessor = keras_hub.models.Gemma3TextEmbedderPreprocessor.from_preset(
+        "harrier_embedding_oss_270m",
+    )
+
+    # Tokenize and pack a single sentence.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize a batch of single sentences.
+    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_hub.models.Gemma3TextEmbedderPreprocessor.from_preset(
+        "harrier_embedding_oss_270m",
+    )
+
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+
+    # Map unlabeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices(first)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    backbone_cls = Gemma3Backbone
+    tokenizer_cls = Gemma3Tokenizer
+
+    def _format_output(self, token_ids, segment_ids):
+        """Format packer output into Gemma3-compatible input dictionary."""
+        return {
+            "token_ids": token_ids,
+            "padding_mask": token_ids != self.tokenizer.pad_token_id,
+        }