keras-hub-nightly 0.23.0.dev202510240418__py3-none-any.whl → 0.24.0.dev202512090431__py3-none-any.whl

keras_hub/src/models/dinov3/dinov3_presets.py

@@ -0,0 +1,93 @@
+"""DINOV3 model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "dinov3_vit_small_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (small-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 21_600_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_small_lvd1689m/1",
+    },
+    "dinov3_vit_small_plus_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (small-plus-sized model) trained on "
+                "LVD-1689M using DINOv3."
+            ),
+            "params": 29_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_small_plus_lvd1689m/1",
+    },
+    "dinov3_vit_base_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (base-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 86_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_base_lvd1689m/1",
+    },
+    "dinov3_vit_large_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (large-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 300_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_large_lvd1689m/1",
+    },
+    "dinov3_vit_huge_plus_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (huge-plus-sized model) trained on "
+                "LVD-1689M using DINOv3."
+            ),
+            "params": 840_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_huge_plus_lvd1689m/1",
+    },
+    "dinov3_vit_7b_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (7B-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 6_700_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_7b_lvd1689m/1",
+    },
+    "dinov3_vit_large_sat493m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (large-sized model) trained on SAT-493M "
+                "using DINOv3."
+            ),
+            "params": 300_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_large_sat493m/1",
+    },
+    "dinov3_vit_7b_sat493m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (7B-sized model) trained on SAT-493M "
+                "using DINOv3."
+            ),
+            "params": 6_700_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_7b_sat493m/1",
+    },
+}

keras_hub/src/models/esm/esm_attention.py

@@ -14,7 +14,8 @@ class ESMRotaryEmbedding(RotaryEmbedding):
         inv_freq = self.scaling_factor / (
             self.max_wavelength ** (ops.arange(0, dim, 2, dtype=x.dtype) / dim)
         )
-        t = ops.arange(x.shape[position], dtype=x.dtype)
+        # Use ops.shape for dynamic shape compatibility with TFLite
+        t = ops.arange(ops.shape(x)[position], dtype=x.dtype)
         freqs = ops.outer(t, inv_freq)
         emb = ops.concatenate((freqs, freqs), axis=-1)
 
@@ -32,11 +33,17 @@ class ESMRotaryEmbedding(RotaryEmbedding):
 
     def rotate_half(self, x):
         x1, x2 = ops.split(x, 2, -1)
-        return ops.concatenate((-x2, x1), axis=-1)
+        # Avoid `ops.concatenate` to prevent XLA compilation issues on JAX
+        # backend. Use stack + reshape approach from base RotaryEmbedding.
+        half_rot_x = ops.stack((-x2, x1), axis=-2)
+        half_rot_x = ops.reshape(half_rot_x, ops.shape(x))
+        return half_rot_x
 
     def apply_rotary_pos_emb(self, x, cos, sin):
-        cos = cos[:, : x.shape[1], :, :]
-        sin = sin[:, : x.shape[1], :, :]
+        # Use ops.shape for dynamic shape compatibility with TFLite
+        seq_len = ops.shape(x)[1]
+        cos = cos[:, :seq_len, :, :]
+        sin = sin[:, :seq_len, :, :]
 
         return (x * cos) + (self.rotate_half(x) * sin)
 

keras_hub/src/models/gemma/gemma_causal_lm.py

@@ -431,3 +431,19 @@ class GemmaCausalLM(CausalLM):
         )
         per_token_loss = per_token_loss_fn(target_ids, logits)
         return per_token_loss
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        # Wrap embedding + scaling
+        backbone = self.backbone
+        inputs = keras.Input(shape=(None,), dtype="int32")
+        x = backbone.token_embedding(inputs)
+        x = x * ops.cast(ops.sqrt(backbone.hidden_dim), x.dtype)
+        pre_processor = keras.Model(inputs=inputs, outputs=x)
+
+        return {
+            "pre_block_layers": [pre_processor],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/gemma3/gemma3_causal_lm_preprocessor.py

@@ -283,9 +283,14 @@ class Gemma3CausalLMPreprocessor(CausalLMPreprocessor):
         # is `None`.
         self.text_only_model = self.image_converter is None
 
-        self.image_placeholder = self.tokenizer.image_placeholder
-        self.start_of_image_token = self.tokenizer.start_of_image_token
-        self.end_of_image_token = self.tokenizer.end_of_image_token
+        if self.text_only_model:
+            self.image_placeholder = None
+            self.start_of_image_token = None
+            self.end_of_image_token = None
+        else:
+            self.image_placeholder = self.tokenizer.image_placeholder
+            self.start_of_image_token = self.tokenizer.start_of_image_token
+            self.end_of_image_token = self.tokenizer.end_of_image_token
 
     def build(self, input_shape):
         # Defer packer creation to `build()` so that we can be sure tokenizer

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -181,4 +181,43 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_270m/4",
     },
+    "medgemma_instruct_4b": {
+        "metadata": {
+            "description": (
+                "A 4 billion parameter model based on Gemma 3. "
+                "This model is trained for performance on medical text"
+                "and image comprehension and is optimized for medical"
+                "applications that involve a text generation component."
+            ),
+            "params": 4300079472,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_4b/1",
+    },
+    "medgemma_instruct_27b": {
+        "metadata": {
+            "description": (
+                "A 27 billion parameter model based on Gemma 3. "
+                "This model trained for performance on medical text "
+                "and image comprehension and is optimized for medical "
+                "applications that involve a text generation component."
+            ),
+            "params": 27432406640,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_27b/1",
+    },
+    "medgemma_instruct_27b_text": {
+        "metadata": {
+            "description": (
+                "A 27 billion parameter text-only model based on Gemma 3. "
+                "This model is trained for performance on medical text "
+                "comprehension and is optimized for medical applications "
+                "that involve a text generation component."
+            ),
+            "params": 27009002240,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_27b_text/1",
+    },
 }

keras_hub/src/models/gemma3/gemma3_tokenizer.py

@@ -77,20 +77,32 @@ class Gemma3Tokenizer(SentencePieceTokenizer):
 
     backbone_cls = Gemma3Backbone
 
-    def __init__(self, proto, **kwargs):
+    def __init__(self, proto, has_vision_tokens=True, **kwargs):
         # Add special tokens.
 
+        self.has_vision_tokens = has_vision_tokens
         # The usual tokens.
         self._add_special_token("<bos>", "start_token")
         self._add_special_token("<eos>", "end_token")
         self._add_special_token("<pad>", "pad_token")
 
-        # Image placeholder token.
-        self._add_special_token("<img>", "image_placeholder")
-
-        #  Some tokens which are used in the preprocessor. We need to keep them
-        # here so that the preprocessor works with `tf.data`.
-        self._add_special_token("<start_of_image>", "start_of_image_token")
-        self._add_special_token("<end_of_image>", "end_of_image_token")
+        if has_vision_tokens:
+            # Image placeholder token.
+            self._add_special_token("<img>", "image_placeholder")
+            # Some tokens which are used in the preprocessor.
+            # We need to keep them
+            # here so that the preprocessor works with tf.data.
+            self._add_special_token("<start_of_image>", "start_of_image_token")
+            self._add_special_token("<end_of_image>", "end_of_image_token")
+        else:
+            # For text-only, skip assigning token IDs or set to -1
+            self.start_of_image_token_id = -1
+            self.image_placeholder_token_id = -1
+            self.end_of_image_token_id = -1
 
         super().__init__(proto=proto, **kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"has_vision_tokens": self.has_vision_tokens})
+        return config

keras_hub/src/models/gpt2/gpt2_causal_lm.py

@@ -420,3 +420,20 @@ class GPT2CausalLM(CausalLM):
         )
         per_token_loss = per_token_loss_fn(target_ids, logits)
         return per_token_loss
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        token_ids = keras.Input(shape=(None,), dtype="int32")
+        tokens = backbone.token_embedding(token_ids)
+        positions = backbone.position_embedding(tokens)
+        x = backbone.embeddings_add((tokens, positions))
+        x = backbone.embeddings_dropout(x)
+        pre_processor = keras.Model(inputs=token_ids, outputs=x)
+
+        return {
+            "pre_block_layers": [pre_processor],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/masked_lm.py

@@ -84,3 +84,25 @@ class MaskedLM(Task):
             weighted_metrics=weighted_metrics,
             **kwargs,
         )
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        # Check for standard backbone structure.
+        if not hasattr(backbone, "transformer_layers"):
+            return None
+
+        # Check for embedding.
+        embedding = getattr(backbone, "token_embedding", None)
+        if embedding is None:
+            embedding = getattr(backbone, "embedding", None)
+
+        if embedding is None:
+            return None
+
+        return {
+            "pre_block_layers": [embedding],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/qwen3/qwen3_presets.py

@@ -70,4 +70,40 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_32b_en/1",
     },
+    "qwen3_embedding_0.6b_en": {
+        "metadata": {
+            "description": (
+                "This text embedding model features a 32k context length and "
+                "offers flexible, user-defined embedding dimensions that can "
+                "range from 32 to 1024."
+            ),
+            "params": 595776512,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3-embedding/keras/qwen3_embedding_0.6b_en/1",
+    },
+    "qwen3_embedding_4b_en": {
+        "metadata": {
+            "description": (
+                "This text embedding model features a 32k context length and "
+                "offers flexible, user-defined embedding dimensions that can "
+                "range from 32 to 2560."
+            ),
+            "params": 4021774336,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3-embedding/keras/qwen3_embedding_4b_en/1",
+    },
+    "qwen3_embedding_8b_en": {
+        "metadata": {
+            "description": (
+                "This text embedding model features a 32k context length and "
+                "offers flexible, user-defined embedding dimensions that can "
+                "range from 32 to 4096."
+            ),
+            "params": 8188515328,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3-embedding/keras/qwen3_embedding_8b_en/1",
+    },
 }

keras_hub/src/models/siglip/siglip_presets.py

@@ -321,4 +321,19 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/siglip/keras/siglip2_so400m_patch16_512/1",
     },
+    "medsiglip_900m_448": {
+        "metadata": {
+            "description": (
+                "A 900 million parameter variant of SigLIP trained to encode "
+                "medical images and text into a common embedding space. "
+                "MedSigLIP contains a vision encoder and a text encoder, and "
+                "supports 448x448 image resolution with up to 64 text tokens."
+            ),
+            "params": 878301426,
+            "official_name": "SigLIP2",
+            "path": "siglip",
+            "model_card": "https://huggingface.co/google/medsiglip-448#medsiglip-model-card",
+        },
+        "kaggle_handle": "kaggle://keras/medsiglip/keras/medsiglip_900m_448/1",
+    },
 }

keras_hub/src/models/smollm3/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.smollm3.smollm3_backbone import SmolLM3Backbone
+from keras_hub.src.models.smollm3.smollm3_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, SmolLM3Backbone)

keras_hub/src/models/smollm3/smollm3_presets.py

@@ -0,0 +1,16 @@
+"""SmolLM3 model preset configurations."""
+
+backbone_presets = {
+    "smollm3_3b_en": {
+        "metadata": {
+            "description": (
+                "Dense decoder-only model has 3 billion total parameters, "
+                "built on 36 layers and utilizes 16 query and "
+                "4 key/value attention heads."
+            ),
+            "params": 3075100928,
+            "path": "smollm3",
+        },
+        "kaggle_handle": "kaggle://keras/smollm3/keras/smollm3_3b_en/1",
+    },
+}

keras_hub/src/utils/tensor_utils.py

@@ -12,9 +12,11 @@ from packaging import version
 
 try:
     import tensorflow as tf
-    import tensorflow_text as tf_text
 except ImportError:
     tf = None
+try:
+    import tensorflow_text as tf_text
+except ImportError:
     tf_text = None
 
 

keras_hub/src/utils/transformers/convert_dinov3.py

@@ -0,0 +1,106 @@
+import numpy as np
+
+from keras_hub.src.models.dinov3.dinov3_backbone import DINOV3Backbone
+
+backbone_cls = DINOV3Backbone
+
+
+def convert_backbone_config(transformers_config):
+    image_size = transformers_config["image_size"]
+    return {
+        "patch_size": transformers_config["patch_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "layer_scale_init_value": transformers_config["layerscale_value"],
+        "num_register_tokens": transformers_config["num_register_tokens"],
+        "use_mask_token": True,
+        "hidden_activation": transformers_config["hidden_act"],
+        "use_gated_mlp": transformers_config["use_gated_mlp"],
+        "use_query_bias": transformers_config["query_bias"],
+        "use_key_bias": transformers_config["key_bias"],
+        "use_value_bias": transformers_config["value_bias"],
+        "use_proj_bias": transformers_config["proj_bias"],
+        "use_mlp_bias": transformers_config["mlp_bias"],
+        "attention_dropout": transformers_config["attention_dropout"],
+        "drop_path_rate": transformers_config["drop_path_rate"],
+        "layer_norm_eps": transformers_config["layer_norm_eps"],
+        "image_shape": (image_size, image_size, 3),
+        "rope_theta": transformers_config["rope_theta"],
+        "apply_layernorm": False,
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    if not isinstance(backbone, DINOV3Backbone):
+        raise ValueError(
+            "The provided backbone must be an instance of DINOV3Backbone. "
+            f"Received: {type(backbone)}"
+        )
+
+    def port_ln(keras_variable, weight_key):
+        loader.port_weight(keras_variable.gamma, f"{weight_key}.weight")
+        loader.port_weight(keras_variable.beta, f"{weight_key}.bias")
+
+    def port_dense(keras_variable, weight_key):
+        loader.port_weight(
+            keras_variable.kernel,
+            f"{weight_key}.weight",
+            hook_fn=lambda x, _: x.T,
+        )
+        if keras_variable.bias is not None:
+            loader.port_weight(keras_variable.bias, f"{weight_key}.bias")
+
+    # Embedding.
+    loader.port_weight(
+        keras_variable=backbone.embeddings.cls_token,
+        hf_weight_key="embeddings.cls_token",
+    )
+    if backbone.use_mask_token:
+        loader.port_weight(
+            keras_variable=backbone.embeddings.mask_token,
+            hf_weight_key="embeddings.mask_token",
+        )
+    if backbone.num_register_tokens > 0:
+        loader.port_weight(
+            keras_variable=backbone.embeddings.register_tokens,
+            hf_weight_key="embeddings.register_tokens",
+        )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.patch_embeddings.projection.kernel,
+        hf_weight_key="embeddings.patch_embeddings.weight",
+        hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.patch_embeddings.projection.bias,
+        hf_weight_key="embeddings.patch_embeddings.bias",
+    )
+
+    # Encoder.
+    for i, layer in enumerate(backbone.encoder.layers):
+        prefix = f"layer.{i}"
+        port_ln(layer.norm1, f"{prefix}.norm1")
+        port_dense(layer.attention.query_dense, f"{prefix}.attention.q_proj")
+        port_dense(layer.attention.key_dense, f"{prefix}.attention.k_proj")
+        port_dense(layer.attention.value_dense, f"{prefix}.attention.v_proj")
+        port_dense(layer.attention.output_dense, f"{prefix}.attention.o_proj")
+
+        loader.port_weight(
+            keras_variable=layer.layer_scale1.lambda1,
+            hf_weight_key=f"{prefix}.layer_scale1.lambda1",
+        )
+        port_ln(layer.norm2, f"{prefix}.norm2")
+        if backbone.use_gated_mlp:
+            port_dense(layer.mlp.gate_proj, f"{prefix}.mlp.gate_proj")
+            port_dense(layer.mlp.up_proj, f"{prefix}.mlp.up_proj")
+            port_dense(layer.mlp.down_proj, f"{prefix}.mlp.down_proj")
+        else:
+            port_dense(layer.mlp.up_proj, f"{prefix}.mlp.up_proj")
+            port_dense(layer.mlp.down_proj, f"{prefix}.mlp.down_proj")
+        loader.port_weight(
+            keras_variable=layer.layer_scale2.lambda1,
+            hf_weight_key=f"{prefix}.layer_scale2.lambda1",
+        )
+
+    port_ln(backbone.layernorm, "norm")