Skip to content

Commit 601b176

Browse files
buildwithsuhanabuildwithsuhana
authored
Standardize rope_theta extraction and enhance converter tests (#2733)
1 parent 154fd48 commit 601b176

38 files changed

Lines changed: 1052 additions & 36 deletions

‎keras_hub/src/utils/transformers/convert_dinov3.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,12 @@
77

88
def convert_backbone_config(transformers_config):
99
image_size = transformers_config["image_size"]
10+
rope_theta = transformers_config.get("rope_parameters", {}).get(
11+
"rope_theta"
12+
)
13+
if rope_theta is None:
14+
rope_theta = transformers_config["rope_theta"]
15+
1016
return {
1117
"patch_size": transformers_config["patch_size"],
1218
"num_layers": transformers_config["num_hidden_layers"],
@@ -27,7 +33,7 @@ def convert_backbone_config(transformers_config):
2733
"drop_path_rate": transformers_config["drop_path_rate"],
2834
"layer_norm_eps": transformers_config["layer_norm_eps"],
2935
"image_shape": (image_size, image_size, 3),
30-
"rope_theta": transformers_config["rope_theta"],
36+
"rope_theta": rope_theta,
3137
"apply_layernorm": False,
3238
}
3339

‎keras_hub/src/utils/transformers/convert_dinov3_test.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
import numpy as np
22
import pytest
33

4+
from keras_hub.src.models.backbone import Backbone
45
from keras_hub.src.models.dinov3.dinov3_backbone import DINOV3Backbone
56
from keras_hub.src.tests.test_case import TestCase
7+
from keras_hub.src.utils.transformers import convert_dinov3
68

79

810
class TestTask(TestCase):
@@ -33,3 +35,67 @@ def test_convert_tiny_preset(self):
3335
],
3436
atol=1e-2,
3537
)
38+
39+
@pytest.mark.extra_large
40+
def test_class_detection(self):
41+
model = Backbone.from_preset(
42+
"hf://facebook/dinov3-vits16-pretrain-lvd1689m",
43+
image_shape=(224, 224, 3),
44+
load_weights=False,
45+
)
46+
self.assertIsInstance(model, DINOV3Backbone)
47+
48+
def test_convert_backbone_config_rope_theta(self):
49+
# transformers < 5 format
50+
transformers_config = {
51+
"image_size": 224,
52+
"patch_size": 16,
53+
"num_hidden_layers": 2,
54+
"hidden_size": 32,
55+
"num_attention_heads": 4,
56+
"intermediate_size": 48,
57+
"layerscale_value": 1.0,
58+
"num_register_tokens": 0,
59+
"hidden_act": "gelu",
60+
"use_gated_mlp": False,
61+
"query_bias": True,
62+
"key_bias": True,
63+
"value_bias": True,
64+
"proj_bias": True,
65+
"mlp_bias": True,
66+
"attention_dropout": 0.0,
67+
"drop_path_rate": 0.0,
68+
"layer_norm_eps": 1e-6,
69+
"rope_theta": 10000.0,
70+
}
71+
keras_config = convert_dinov3.convert_backbone_config(
72+
transformers_config
73+
)
74+
self.assertEqual(keras_config["rope_theta"], 10000.0)
75+
76+
# transformers >= 5 format
77+
transformers_config = {
78+
"image_size": 224,
79+
"patch_size": 16,
80+
"num_hidden_layers": 2,
81+
"hidden_size": 32,
82+
"num_attention_heads": 4,
83+
"intermediate_size": 48,
84+
"layerscale_value": 1.0,
85+
"num_register_tokens": 0,
86+
"hidden_act": "gelu",
87+
"use_gated_mlp": False,
88+
"query_bias": True,
89+
"key_bias": True,
90+
"value_bias": True,
91+
"proj_bias": True,
92+
"mlp_bias": True,
93+
"attention_dropout": 0.0,
94+
"drop_path_rate": 0.0,
95+
"layer_norm_eps": 1e-6,
96+
"rope_parameters": {"rope_theta": 20000.0},
97+
}
98+
keras_config = convert_dinov3.convert_backbone_config(
99+
transformers_config
100+
)
101+
self.assertEqual(keras_config["rope_theta"], 20000.0)

‎keras_hub/src/utils/transformers/convert_gemma3.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,9 @@ def convert_backbone_config(transformers_config):
5656
# global rotary embedding. `rope_parameters` is optional and not used
5757
# by HF for global scaling when `rope_scaling` is None.
5858
rope_scaling = transformer_config.get("rope_scaling", None)
59-
rope_params = transformer_config.get("rope_parameters") or {}
59+
rope_params = transformer_config.get("rope_parameters", {})
60+
if rope_params is None:
61+
rope_params = {}
6062

6163
if rope_scaling is not None:
6264
rope_global_config = rope_scaling or {}

‎keras_hub/src/utils/transformers/convert_gemma3n.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,10 @@ def convert_backbone_config(transformers_config):
181181
"hidden_activation", "gelu_approximate"
182182
)
183183

184+
rope_theta = text_config.get("rope_parameters", {}).get("rope_theta")
185+
if rope_theta is None:
186+
rope_theta = text_config["rope_theta"]
187+
184188
return {
185189
"text_vocab_size": text_config["vocab_size"],
186190
"text_hidden_size": text_config["hidden_size"],
@@ -193,7 +197,7 @@ def convert_backbone_config(transformers_config):
193197
"hidden_activation": hidden_activation,
194198
"layer_types": text_config["layer_types"],
195199
"sliding_window": text_config["sliding_window"],
196-
"rope_theta": text_config["rope_theta"],
200+
"rope_theta": rope_theta,
197201
"max_position_embeddings": text_config["max_position_embeddings"],
198202
"vocab_size_per_layer_input": text_config["vocab_size_per_layer_input"],
199203
"hidden_size_per_layer_input": text_config[

‎keras_hub/src/utils/transformers/convert_gemma3n_test.py‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
import pytest
2+
3+
from keras_hub.src.models.backbone import Backbone
4+
from keras_hub.src.models.gemma3n.gemma3n_backbone import Gemma3nBackbone
5+
from keras_hub.src.tests.test_case import TestCase
6+
from keras_hub.src.utils.transformers import convert_gemma3n
7+
8+
9+
class TestGemma3nConverter(TestCase):
10+
@pytest.mark.extra_large
11+
def test_convert_tiny_preset(self):
12+
model = Gemma3nBackbone.from_preset(
13+
"hf://yujiepan/gemma-3n-tiny-random-dim4",
14+
load_weights=False,
15+
)
16+
self.assertIsInstance(model, Gemma3nBackbone)
17+
18+
@pytest.mark.large
19+
def test_class_detection(self):
20+
model = Backbone.from_preset(
21+
"hf://yujiepan/gemma-3n-tiny-random-dim4",
22+
load_weights=False,
23+
)
24+
self.assertIsInstance(model, Gemma3nBackbone)
25+
26+
def test_gemma3n_rope_theta(self):
27+
# transformers < 5 format
28+
text_cfg = {
29+
"vocab_size": 100,
30+
"hidden_size": 32,
31+
"num_hidden_layers": 2,
32+
"num_attention_heads": 4,
33+
"num_key_value_heads": 2,
34+
"head_dim": 8,
35+
"intermediate_size": 48,
36+
"layer_types": ["full_attention", "full_attention"],
37+
"sliding_window": 4096,
38+
"rope_theta": 10000.0,
39+
"rms_norm_eps": 1e-5,
40+
"max_position_embeddings": 32,
41+
"vocab_size_per_layer_input": 100,
42+
"hidden_size_per_layer_input": 32,
43+
"altup_num_inputs": 1,
44+
"laurel_rank": 1,
45+
"attention_bias": False,
46+
"attention_dropout": 0.0,
47+
"altup_coef_clip": 1.0,
48+
"altup_active_idx": 0,
49+
"altup_correct_scale": False,
50+
"num_kv_shared_layers": 1,
51+
}
52+
transformers_config = {
53+
"text_config": text_cfg,
54+
"vision_config": {
55+
"hidden_size": 32,
56+
"vocab_size": 100,
57+
"vocab_offset": 0,
58+
},
59+
"audio_config": {
60+
"hidden_size": 32,
61+
"input_feat_size": 32,
62+
"sscp_conv_channel_size": 32,
63+
"sscp_conv_kernel_size": 3,
64+
"sscp_conv_stride_size": 2,
65+
"sscp_conv_group_norm_eps": 1e-5,
66+
"conf_num_hidden_layers": 2,
67+
"rms_norm_eps": 1e-5,
68+
"gradient_clipping": 1.0,
69+
"conf_residual_weight": 0.1,
70+
"conf_num_attention_heads": 4,
71+
"conf_attention_chunk_size": 16,
72+
"conf_attention_context_right": 8,
73+
"conf_attention_context_left": 8,
74+
"conf_attention_logit_cap": 10.0,
75+
"conf_conv_kernel_size": 3,
76+
"conf_reduction_factor": 2,
77+
"vocab_size": 100,
78+
"vocab_offset": 0,
79+
},
80+
"vision_soft_tokens_per_image": 16,
81+
"image_token_id": 1,
82+
"audio_soft_tokens_per_image": 16,
83+
"audio_token_id": 2,
84+
}
85+
keras_config = convert_gemma3n.convert_backbone_config(
86+
transformers_config
87+
)
88+
self.assertEqual(keras_config["rope_theta"], 10000.0)
89+
90+
# transformers >= 5 format
91+
text_cfg["rope_parameters"] = {"rope_theta": 20000.0}
92+
keras_config = convert_gemma3n.convert_backbone_config(
93+
transformers_config
94+
)
95+
self.assertEqual(keras_config["rope_theta"], 20000.0)

‎keras_hub/src/utils/transformers/convert_gemma4.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,9 @@ def convert_backbone_config(transformers_config):
231231

232232
# Partial RoPE factor for global (full) attention layers.
233233
# Stored under rope_parameters["full_attention"]["partial_rotary_factor"].
234-
rope_params = text_cfg.get("rope_parameters", {}) or {}
234+
rope_params = text_cfg.get("rope_parameters", {})
235+
if rope_params is None:
236+
rope_params = {}
235237
global_rope_partial_rotary_factor = rope_params.get(
236238
"full_attention", {}
237239
).get("partial_rotary_factor")
@@ -243,6 +245,8 @@ def convert_backbone_config(transformers_config):
243245
# If it's missing in `full_attention`, safely fall back to top-level cfg.
244246
if global_rope_theta is None:
245247
global_rope_theta = text_cfg.get("rope_theta")
248+
if local_rope_theta is None:
249+
local_rope_theta = text_cfg.get("rope_theta")
246250

247251
# HF `use_bidirectional_attention` controls vision-token attention only:
248252
# null → purely causal for all tokens (E2B, E4B).

‎keras_hub/src/utils/transformers/convert_gpt_oss.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@
1212

1313
def convert_backbone_config(transformers_config):
1414
"""Convert a Hugging Face Gpt-Oss config to a KerasHub config."""
15+
rope_theta = transformers_config.get("rope_parameters", {}).get(
16+
"rope_theta"
17+
)
18+
if rope_theta is None:
19+
rope_theta = transformers_config["rope_theta"]
1520
config = {
1621
"vocabulary_size": transformers_config["vocab_size"],
1722
"num_layers": transformers_config["num_hidden_layers"],
@@ -21,7 +26,7 @@ def convert_backbone_config(transformers_config):
2126
"num_key_value_heads": transformers_config["num_key_value_heads"],
2227
"num_experts": transformers_config["num_local_experts"],
2328
"top_k": transformers_config["num_experts_per_tok"],
24-
"rope_max_wavelength": transformers_config["rope_theta"],
29+
"rope_max_wavelength": rope_theta,
2530
"layer_norm_epsilon": transformers_config["rms_norm_eps"],
2631
"sliding_window": transformers_config.get("sliding_window"),
2732
"output_router_logits": transformers_config.get(

‎keras_hub/src/utils/transformers/convert_gpt_oss_test.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from keras_hub.src.models.gpt_oss.gpt_oss_backbone import GptOssBackbone
66
from keras_hub.src.models.gpt_oss.gpt_oss_causal_lm import GptOssCausalLM
77
from keras_hub.src.tests.test_case import TestCase
8+
from keras_hub.src.utils.transformers import convert_gpt_oss
89

910

1011
class TestTask(TestCase):
@@ -27,3 +28,42 @@ def test_class_detection(self):
2728
load_weights=False,
2829
)
2930
self.assertIsInstance(model, GptOssBackbone)
31+
32+
def test_convert_backbone_config_rope_theta(self):
33+
# transformers < 5 format
34+
transformers_config = {
35+
"vocab_size": 100,
36+
"num_hidden_layers": 2,
37+
"num_attention_heads": 4,
38+
"hidden_size": 32,
39+
"intermediate_size": 48,
40+
"num_key_value_heads": 2,
41+
"num_local_experts": 2,
42+
"num_experts_per_tok": 1,
43+
"rope_theta": 10000.0,
44+
"rms_norm_eps": 1e-5,
45+
"sliding_window": 4096,
46+
}
47+
keras_config = convert_gpt_oss.convert_backbone_config(
48+
transformers_config
49+
)
50+
self.assertEqual(keras_config["rope_max_wavelength"], 10000.0)
51+
52+
# transformers >= 5 format
53+
transformers_config = {
54+
"vocab_size": 100,
55+
"num_hidden_layers": 2,
56+
"num_attention_heads": 4,
57+
"hidden_size": 32,
58+
"intermediate_size": 48,
59+
"num_key_value_heads": 2,
60+
"num_local_experts": 2,
61+
"num_experts_per_tok": 1,
62+
"rope_parameters": {"rope_theta": 20000.0},
63+
"rms_norm_eps": 1e-5,
64+
"sliding_window": 4096,
65+
}
66+
keras_config = convert_gpt_oss.convert_backbone_config(
67+
transformers_config
68+
)
69+
self.assertEqual(keras_config["rope_max_wavelength"], 20000.0)

‎keras_hub/src/utils/transformers/convert_llama3.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,11 @@
77

88

99
def convert_backbone_config(transformers_config):
10+
rope_theta = transformers_config.get("rope_parameters", {}).get(
11+
"rope_theta"
12+
)
13+
if rope_theta is None:
14+
rope_theta = transformers_config["rope_theta"]
1015
backbone_config = {
1116
"vocabulary_size": transformers_config["vocab_size"],
1217
"num_layers": transformers_config["num_hidden_layers"],
@@ -15,7 +20,7 @@ def convert_backbone_config(transformers_config):
1520
"intermediate_dim": transformers_config["intermediate_size"],
1621
"num_key_value_heads": transformers_config["num_key_value_heads"],
1722
"tie_word_embeddings": transformers_config["tie_word_embeddings"],
18-
"rope_max_wavelength": transformers_config["rope_theta"],
23+
"rope_max_wavelength": rope_theta,
1924
}
2025

2126
if transformers_config.get("rope_scaling", None) is not None:

‎keras_hub/src/utils/transformers/convert_llama3_test.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from keras_hub.src.models.llama3.llama3_backbone import Llama3Backbone
66
from keras_hub.src.models.llama3.llama3_causal_lm import Llama3CausalLM
77
from keras_hub.src.tests.test_case import TestCase
8+
from keras_hub.src.utils.transformers import convert_llama3
89

910

1011
class TestTask(TestCase):
@@ -27,4 +28,39 @@ def test_class_detection(self):
2728
)
2829
self.assertIsInstance(model, Llama3Backbone)
2930

31+
def test_convert_backbone_config_rope_theta(self):
32+
# transformers < 5 format
33+
transformers_config = {
34+
"vocab_size": 100,
35+
"num_hidden_layers": 2,
36+
"num_attention_heads": 4,
37+
"hidden_size": 32,
38+
"intermediate_size": 48,
39+
"num_key_value_heads": 2,
40+
"rope_theta": 10000.0,
41+
"rms_norm_eps": 1e-5,
42+
"tie_word_embeddings": True,
43+
}
44+
keras_config = convert_llama3.convert_backbone_config(
45+
transformers_config
46+
)
47+
self.assertEqual(keras_config["rope_max_wavelength"], 10000.0)
48+
49+
# transformers >= 5 format
50+
transformers_config = {
51+
"vocab_size": 100,
52+
"num_hidden_layers": 2,
53+
"num_attention_heads": 4,
54+
"hidden_size": 32,
55+
"intermediate_size": 48,
56+
"num_key_value_heads": 2,
57+
"rope_parameters": {"rope_theta": 20000.0},
58+
"rms_norm_eps": 1e-5,
59+
"tie_word_embeddings": True,
60+
}
61+
keras_config = convert_llama3.convert_backbone_config(
62+
transformers_config
63+
)
64+
self.assertEqual(keras_config["rope_max_wavelength"], 20000.0)
65+
3066
# TODO: compare numerics with huggingface model

0 commit comments

Comments
 (0)