keras-hub-nightly 0.23.0.dev202508300415__tar.gz → 0.23.0.dev202509130423__tar.gz

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.23.0.dev202508300415
+Version: 0.23.0.dev202509130423
 Summary: Pretrained models for Keras.
 Author-email: Keras team <keras-users@googlegroups.com>
 License-Expression: Apache-2.0

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/layers/modeling/reversible_embedding.py

@@ -235,7 +235,8 @@ class ReversibleEmbedding(keras.layers.Embedding):
 
         return super()._int8_call(inputs)
 
-    def quantize(self, mode, type_check=True):
+    def quantize(self, mode, type_check=True, config=None):
+        del config
         if type_check and type(self) is not ReversibleEmbedding:
             raise self._not_implemented_error(self.quantize)
 

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/backbone.py

@@ -130,7 +130,8 @@ class Backbone(keras.Model):
         1. a built-in preset identifier like `'bert_base_en'`
         2. a Kaggle Models handle like `'kaggle://user/bert/keras/bert_base_en'`
         3. a Hugging Face handle like `'hf://user/bert_base_en'`
-        4. a path to a local preset directory like `'./bert_base_en'`
+        4. a ModelScope handle like `'modelscope://user/bert_base_en'`
+        5. a path to a local preset directory like `'./bert_base_en'`
 
         This constructor can be called in one of two ways. Either from the base
         class like `keras_hub.models.Backbone.from_preset()`, or from

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/causal_lm.py

@@ -132,6 +132,17 @@ class CausalLM(Task):
             return self.generate_function
 
         self.generate_function = self.generate_step
+        if keras.config.backend() == "openvino":
+            from keras_hub.src.utils.openvino_utils import ov_infer
+
+            def wrapped_generate_function(inputs, stop_token_ids=None):
+                # Convert to numpy for OpenVINO backend
+                inputs = tree.map_structure(ops.array, inputs)
+                return ov_infer(
+                    self, inputs, stop_token_ids, self.generate_step
+                )
+
+            self.generate_function = wrapped_generate_function
         if keras.config.backend() == "torch":
             import torch
 

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/gemma/gemma_presets.py

@@ -198,4 +198,12 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_27b_en/2",
     },
+    "vault_gemma_1b_en": {
+        "metadata": {
+            "description": "1 billion parameter, 26-layer, VaultGemma model.",
+            "params": 1038741120,
+            "path": "gemma",
+        },
+        "kaggle_handle": "kaggle://keras/vaultgemma/keras/vault_gemma_1b_en/2",
+    },
 }

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/gemma3/gemma3_attention.py

@@ -46,6 +46,7 @@ class CachedGemma3Attention(keras.layers.Layer):
         layer_norm_epsilon=1e-6,
         rope_wavelength=10_000.0,
         rope_scaling_factor=1.0,
+        use_bidirectional_attention=False,
         dropout=0,
         **kwargs,
     ):
@@ -61,6 +62,7 @@ class CachedGemma3Attention(keras.layers.Layer):
         self.layer_norm_epsilon = layer_norm_epsilon
         self.rope_wavelength = rope_wavelength
         self.rope_scaling_factor = rope_scaling_factor
+        self.use_bidirectional_attention = use_bidirectional_attention
         self.dropout = dropout
 
         self._kernel_initializer = keras.initializers.get(
@@ -240,12 +242,58 @@ class CachedGemma3Attention(keras.layers.Layer):
         results = ops.einsum("bkgts,bskh->btkgh", attention_softmax, v)
         return ops.reshape(results, (b, q_len, self.num_query_heads, h))
 
+    def _compute_bidirectional_sliding_mask(self, batch_size, sequence_length):
+        """Computes a bidirectional sliding window attention mask.
+
+        A token can attend to any other token if their absolute distance is
+        within  half the sliding window size. This mask is used in embedding
+        models like `EmbeddingGemma`.
+
+        Args:
+            batch_size: The batch size for the mask.
+            sequence_length: The length of the sequence.
+
+        Returns:
+            A boolean attention mask with shape
+            `(batch_size, sequence_length, sequence_length)`.
+        """
+        i = keras.ops.expand_dims(
+            keras.ops.arange(sequence_length, dtype="int32"), axis=1
+        )
+        j = keras.ops.arange(sequence_length, dtype="int32")
+
+        # If sliding window size is 4, the token in question attends to 1
+        # token before and 2 tokens after.
+        w_right = self.sliding_window_size // 2
+        w_left = self.sliding_window_size - w_right - 1
+
+        # Calculate the relative distance.
+        distance = i - j
+
+        mask = keras.ops.logical_and(distance <= w_left, distance >= -w_right)
+
+        mask = keras.ops.expand_dims(mask, axis=0)
+        return keras.ops.broadcast_to(
+            mask, (batch_size, sequence_length, sequence_length)
+        )
+
     def _mask_sliding_window(
         self,
         attention_mask,
         cache_update_index=0,
     ):
         batch_size, query_len, key_len = ops.shape(attention_mask)
+
+        if self.use_bidirectional_attention:
+            bidirectional_sliding_mask = (
+                self._compute_bidirectional_sliding_mask(
+                    batch_size=batch_size,
+                    # `query_len = key_len` for embedding models
+                    sequence_length=query_len,
+                )
+            )
+            return ops.logical_and(attention_mask, bidirectional_sliding_mask)
+
         # Compute the sliding window for square attention.
         all_ones = ops.ones((key_len, key_len), "bool")
         if keras.config.backend() == "tensorflow":

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/gemma3/gemma3_backbone.py

@@ -196,6 +196,7 @@ class Gemma3Backbone(Backbone):
         global_rope_scaling_factor=1.0,
         vision_encoder=None,
         layer_norm_epsilon=1e-6,
+        use_bidirectional_attention=False,
         dropout=0,
         dtype=None,
         **kwargs,
@@ -251,6 +252,7 @@ class Gemma3Backbone(Backbone):
                 sliding_window_size=sliding_window_size,
                 rope_wavelength=rope_wavelength,
                 rope_scaling_factor=rope_scaling_factor,
+                use_bidirectional_attention=use_bidirectional_attention,
                 dropout=dropout,
                 dtype=dtype,
                 name=f"decoder_block_{i}",
@@ -357,6 +359,7 @@ class Gemma3Backbone(Backbone):
         self.sliding_window_size = sliding_window_size
         self.local_rope_scaling_factor = local_rope_scaling_factor
         self.global_rope_scaling_factor = global_rope_scaling_factor
+        self.use_bidirectional_attention = use_bidirectional_attention
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
 
@@ -396,6 +399,7 @@ class Gemma3Backbone(Backbone):
                 "vision_encoder": None
                 if self.vision_encoder is None
                 else keras.layers.serialize(self.vision_encoder),
+                "use_bidirectional_attention": self.use_bidirectional_attention,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,
             }

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/gemma3/gemma3_decoder_block.py

@@ -45,6 +45,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
         layer_norm_epsilon=1e-6,
         rope_wavelength=10_000.0,
         rope_scaling_factor=1.0,
+        use_bidirectional_attention=False,
         dropout=0,
         **kwargs,
     ):
@@ -66,6 +67,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
         self.layer_norm_epsilon = layer_norm_epsilon
         self.rope_wavelength = rope_wavelength
         self.rope_scaling_factor = rope_scaling_factor
+        self.use_bidirectional_attention = use_bidirectional_attention
         self.dropout = dropout
 
         self.pre_attention_norm = RMSNormalization(
@@ -93,6 +95,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
             rope_wavelength=rope_wavelength,
             rope_scaling_factor=rope_scaling_factor,
             dropout=dropout,
+            use_bidirectional_attention=use_bidirectional_attention,
             dtype=self.dtype_policy,
             name="attention",
         )
@@ -209,6 +212,14 @@ class Gemma3DecoderBlock(keras.layers.Layer):
         if cache is not None:
             input_length = ops.shape(cache)[2]
 
+        if self.use_bidirectional_attention:
+            # `output_length` and `input_length` will be the same in this case
+            # because we use bidirectional attention for models like
+            # `EmbeddingGemma` which aren't used for text generation.
+            mask_1 = decoder_mask
+            mask_2 = ops.transpose(mask_1, (0, 2, 1))
+            return mask_1 * mask_2
+
         causal_mask = compute_causal_mask(
             batch_size=batch_size,
             input_length=input_length,
@@ -304,6 +315,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
                 "dropout": self.dropout,
                 "rope_wavelength": self.rope_wavelength,
                 "rope_scaling_factor": self.rope_scaling_factor,
+                "use_bidirectional_attention": self.use_bidirectional_attention,
             }
         )
         return config

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/models/sam/sam_prompt_encoder.py

@@ -67,7 +67,9 @@ class SAMPromptEncoder(keras.layers.Layer):
         self.activation = activation
 
         self.positional_embedding_layer = RandomFrequencyPositionalEmbeddings(
-            num_positional_features=self.hidden_size // 2, scale=1
+            num_positional_features=self.hidden_size // 2,
+            scale=1,
+            dtype=self.dtype,
         )
 
         self.foreground_point_embed = keras.layers.Embedding(

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/samplers/beam_sampler.py

@@ -95,15 +95,15 @@ class BeamSampler(Sampler):
         )
         log_probs = flatten_beams(ops.repeat(log_probs, batch_size, axis=0))
 
-        def cond(prompt, cache, index, log_probs):
+        def cond(prompt, cache, index, mask, log_probs):
             if stop_token_ids is None:
-                return True
+                return ops.convert_to_tensor(True, dtype="bool")
             # Stop if all sequences have produced a *new* stop token.
             end_tokens = any_equal(prompt, stop_token_ids, ~mask)
             prompt_done = ops.any(end_tokens, axis=-1)
             return ops.logical_not(ops.all(prompt_done))
 
-        def body(prompt, cache, index, log_probs):
+        def body(prompt, cache, index, mask, log_probs):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
             vocab_size = ops.shape(logits)[-1]
@@ -150,12 +150,12 @@ class BeamSampler(Sampler):
             next_token = next_token[:, None]
             prompt = ops.slice_update(prompt, [0, index], next_token)
             # Return the iteration of the loop state.
-            return (prompt, cache, index + 1, log_probs)
+            return (prompt, cache, index + 1, mask, log_probs)
 
-        prompt, _, _, log_probs = self.run_loop(
+        prompt, _, _, _, log_probs = self.run_loop(
             cond=cond,
             body=body,
-            loop_vars=(prompt, cache, index, log_probs),
+            loop_vars=(prompt, cache, index, mask, log_probs),
             maximum_iterations=(max_length - index),
             model=model,
         )

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/samplers/sampler.py

@@ -92,16 +92,18 @@ class Sampler:
         # `ops.while_loop` will not accept `None` as a value for `loop_vars`.
         cache = () if cache is None else cache
 
-        def cond(prompt, cache, index):
+        # OpenVINO requires all parameters to be passed in the body.
+        # So we pass `mask` as well.
+        def cond(prompt, cache, index, mask):
             if stop_token_ids is None:
-                return True
+                return ops.convert_to_tensor(True, dtype="bool")
             # Stop if all sequences have produced a *new* id from
             # stop_token_ids.
             end_tokens = any_equal(prompt, stop_token_ids, ~mask)
             prompt_done = ops.any(end_tokens, axis=-1)
             return ops.logical_not(ops.all(prompt_done))
 
-        def body(prompt, cache, index):
+        def body(prompt, cache, index, mask):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
             probabilities = self.compute_probabilities(logits)
@@ -115,12 +117,12 @@ class Sampler:
             prompt = ops.slice_update(prompt, [0, index], next_token)
 
             # Return the next prompt, cache and incremented index.
-            return (prompt, cache, index + 1)
+            return (prompt, cache, index + 1, mask)
 
-        prompt, _, _ = self.run_loop(
+        prompt, _, _, _ = self.run_loop(
             cond,
             body,
-            loop_vars=(prompt, cache, index),
+            loop_vars=(prompt, cache, index, mask),
             maximum_iterations=(max_length - index),
             model=model,
         )

keras_hub_nightly-0.23.0.dev202509130423/keras_hub/src/utils/openvino_utils.py

@@ -0,0 +1,141 @@
+from keras import tree
+
+from keras_hub.src.utils.keras_utils import print_msg
+
+try:
+    import openvino as ov
+    import openvino.opset14 as ov_opset
+    from openvino import Core
+except ImportError:
+    ov = None
+    ov_opset = None
+    Core = None
+
+
+_core = None
+
+
+def get_core():
+    """Get or create OpenVINO Core instance.
+
+    Returns:
+        openvino.Core: OpenVINO Core instance,
+        or None if OpenVINO not available.
+    """
+    global _core
+    if _core is None and Core is not None:
+        _core = Core()
+    return _core
+
+
+def get_device():
+    """Detect and return the best available OpenVINO device.
+
+    Returns:
+        str: "GPU" if available, otherwise "CPU".
+    """
+    core = get_core()
+    if core is None:
+        return "CPU"
+    return "GPU" if "GPU" in core.available_devices else "CPU"
+
+
+def compile_model(struct_params, struct_outputs, device, model_dtype):
+    """Compile OpenVINO model with dynamic shapes and precision hints.
+
+    Args:
+        struct_params: Model parameters structure.
+        struct_outputs: Model outputs structure.
+        device: Target device ("GPU" or "CPU").
+        model_dtype: Model precision ("f16" or "f32").
+
+    Returns:
+        Compiled OpenVINO model ready for inference.
+    """
+    flat_params = tree.flatten(struct_params)
+    flat_outputs = tree.flatten(struct_outputs)
+    parameters = [p.output.get_node() for p in flat_params]
+    results = [ov_opset.result(r.output) for r in flat_outputs]
+    ov_model = ov.Model(results=results, parameters=parameters)
+    for ov_input in ov_model.inputs:
+        rank = ov_input.get_partial_shape().rank.get_length()
+        ov_input.get_node().set_partial_shape(ov.PartialShape([-1] * rank))
+    ov_model.validate_nodes_and_infer_types()
+    config = {"INFERENCE_PRECISION_HINT": model_dtype}
+    core = get_core()
+    if core is None:
+        raise RuntimeError("OpenVINO not available")
+    return core.compile_model(ov_model, device, config)
+
+
+def get_outputs(inputs, struct_outputs, compiled_ov_model, unpack_singleton):
+    """Execute compiled OpenVINO model and return structured outputs.
+
+    Args:
+        inputs: Input tensors for inference.
+        struct_outputs: Expected output structure.
+        compiled_ov_model: Compiled OpenVINO model.
+        unpack_singleton: Function to unpack singleton outputs.
+
+    Returns:
+        Structured model outputs matching expected format.
+    """
+    flatten_inputs = tree.flatten(inputs)
+    raw = compiled_ov_model(flatten_inputs).to_tuple()
+    packed = tree.pack_sequence_as(struct_outputs, raw)
+    return unpack_singleton(packed)
+
+
+def ov_infer(model, inputs, stop_token_ids, fn):
+    """High-level OpenVINO inference with model reuse and compilation.
+
+    This function manages OpenVINO model compilation and caching. It reuses
+    existing compiled models when possible, or compiles new ones as needed.
+    Handles device detection and automatic precision selection.
+
+    Args:
+        model: Keras model with OpenVINO backend support.
+        inputs: Input tensors for inference.
+        stop_token_ids: Token IDs that should stop generation.
+        fn: Function to execute with the parameterized inputs.
+
+    Returns:
+        Model outputs from OpenVINO inference.
+    """
+    device = get_device()
+
+    # Try to use existing compiled model for the same device
+    if (
+        getattr(model, "ov_compiled_model", None) is not None
+        and getattr(model, "ov_device", None) is not None
+        and device == model.ov_device
+    ):
+        try:
+            return get_outputs(
+                inputs,
+                model.struct_outputs,
+                model.ov_compiled_model,
+                model._unpack_singleton,
+            )
+        except RuntimeError as e:
+            print_msg(
+                "WARNING: OpenVINO inference \033[1mFAILED\033[0m, "
+                "recompiling model and trying again.\n" + str(e)
+            )
+            model.ov_compiled_model = None
+            model.struct_outputs = None
+
+    # Compile a new model
+    struct_params = model._parameterize_data(inputs)
+    model.struct_outputs = fn(struct_params, stop_token_ids)
+    model.ov_device = device
+    model_dtype = "f16" if model.dtype in ("float16", "bfloat16") else "f32"
+    model.ov_compiled_model = compile_model(
+        struct_params, model.struct_outputs, device, model_dtype
+    )
+    return get_outputs(
+        inputs,
+        model.struct_outputs,
+        model.ov_compiled_model,
+        model._unpack_singleton,
+    )

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub/src/version.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.23.0.dev202508300415"
+__version__ = "0.23.0.dev202509130423"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.23.0.dev202508300415
+Version: 0.23.0.dev202509130423
 Summary: Pretrained models for Keras.
 Author-email: Keras team <keras-users@googlegroups.com>
 License-Expression: Apache-2.0

{keras_hub_nightly-0.23.0.dev202508300415 → keras_hub_nightly-0.23.0.dev202509130423}/keras_hub_nightly.egg-info/SOURCES.txt

@@ -540,6 +540,7 @@ keras_hub/src/tokenizers/word_piece_tokenizer.py
 keras_hub/src/tokenizers/word_piece_tokenizer_trainer.py
 keras_hub/src/utils/__init__.py
 keras_hub/src/utils/keras_utils.py
+keras_hub/src/utils/openvino_utils.py
 keras_hub/src/utils/pipeline_model.py
 keras_hub/src/utils/preset_utils.py
 keras_hub/src/utils/python_utils.py