From e6ba5b2d3c3ec40208b0469c427ceba2c6281507 Mon Sep 17 00:00:00 2001 From: Jyotinder Singh <33001894+JyotinderSingh@users.noreply.github.com> Date: Mon, 1 Sep 2025 13:34:09 +0530 Subject: [PATCH 1/5] Adds support for gemma_270m to checkpoint converter --- .../convert_gemma3_checkpoints.py | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py index 19945cb6ee..6a8e5214dc 100644 --- a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py +++ b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py @@ -10,8 +10,8 @@ Usage: ```shell cd tools/checkpoint_conversion -python convert_gemma_checkpoints.py --preset gemma3_instruct_1b -python convert_gemma_checkpoints.py --preset gemma3_instruct_4b +python convert_gemma3_checkpoints.py --preset gemma3_instruct_1b +python convert_gemma3_checkpoints.py --preset gemma3_instruct_4b ``` """ @@ -43,6 +43,15 @@ PRESET_MAP = { # === Text === + # 270M + "gemma3_instruct_270m": { + "model": gm.nn.Gemma3_270M, + "params": gm.ckpts.CheckpointPath.GEMMA3_270M_IT + }, + "gemma3_270m": { + "model": gm.nn.Gemma3_270M, + "params": gm.ckpts.CheckpointPath.GEMMA3_270M_PT + }, # 1B "gemma3_1b": { "model": gm.nn.Gemma3_1B, @@ -493,11 +502,15 @@ def validate_output( params=flax_params, multi_turn=False, cache_length=256 if length <= 256 else 512, - # max_out_length=length, ) flax_output = flax_sampler.chat(input_str, images=image) print("🔶 Flax output:", flax_output) + if flax_output.startswith(keras_output): + print("✅ Output validated!") + else: + print("❌ Output does not match!") + def main(_): preset = FLAGS.preset @@ -508,11 +521,11 @@ def main(_): assert preset in presets, ( f"Invalid preset {preset}. Must be one of {','.join(presets)}" ) - text_only = "text" in preset or "1b" in preset + text_only = "text" in preset or "1b" in preset or "270m" in preset print("🏃 Loading Flax model and tokeniser") flax_kwargs = {} - if text_only and "1b" not in preset: + if text_only and "1b" not in preset and "270m" not in preset: flax_kwargs["text_only"] = True flax_model = PRESET_MAP[preset]["model"](**flax_kwargs) flax_config = flax_model.config @@ -543,7 +556,7 @@ def main(_): keras_image_converter, flax_model, flax_params, - text_only, + text_only ) keras_model.save_to_preset(preset) From 451f6a56a21c39a09e8a185aa63adf685192bb08 Mon Sep 17 00:00:00 2001 From: Jyotinder Singh <33001894+JyotinderSingh@users.noreply.github.com> Date: Mon, 1 Sep 2025 13:58:07 +0530 Subject: [PATCH 2/5] format --- tools/checkpoint_conversion/convert_gemma3_checkpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py index 6a8e5214dc..1cbdbe1c4c 100644 --- a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py +++ b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py @@ -556,7 +556,7 @@ def main(_): keras_image_converter, flax_model, flax_params, - text_only + text_only, ) keras_model.save_to_preset(preset) From 082ae31c960acb14ae4f0750ec52d53e9d89cfa2 Mon Sep 17 00:00:00 2001 From: Jyotinder Singh <33001894+JyotinderSingh@users.noreply.github.com> Date: Mon, 8 Sep 2025 11:56:30 +0530 Subject: [PATCH 3/5] removes assertion --- tools/checkpoint_conversion/convert_gemma3_checkpoints.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py index 1cbdbe1c4c..d17c912446 100644 --- a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py +++ b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py @@ -506,11 +506,6 @@ def validate_output( flax_output = flax_sampler.chat(input_str, images=image) print("🔶 Flax output:", flax_output) - if flax_output.startswith(keras_output): - print("✅ Output validated!") - else: - print("❌ Output does not match!") - def main(_): preset = FLAGS.preset From 53646be0bea497a8b4a7ca59d231d07f154cb2bb Mon Sep 17 00:00:00 2001 From: Jyotinder Singh <33001894+JyotinderSingh@users.noreply.github.com> Date: Tue, 9 Sep 2025 11:31:17 +0530 Subject: [PATCH 4/5] format --- .vscode/launch.json | 29 +++++++++++++++++++ .../convert_gemma3_checkpoints.py | 4 +-- 2 files changed, 31 insertions(+), 2 deletions(-) create mode 100644 .vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000000..8659be2f04 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,29 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "Pytest: Prompt for test target", + "type": "debugpy", + "request": "launch", + "module": "pytest", + "console": "integratedTerminal", + "justMyCode": true, + "args": [ + "${input:pytestTarget}" + ] + } + ], + + // 2. This 'inputs' array defines the prompt + "inputs": [ + { + "id": "pytestTarget", + "type": "promptString", + "description": "Enter the pytest target (file path or file::function)", + "default": "" + } + ] +} \ No newline at end of file diff --git a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py index d17c912446..2105ae4ff1 100644 --- a/tools/checkpoint_conversion/convert_gemma3_checkpoints.py +++ b/tools/checkpoint_conversion/convert_gemma3_checkpoints.py @@ -46,11 +46,11 @@ # 270M "gemma3_instruct_270m": { "model": gm.nn.Gemma3_270M, - "params": gm.ckpts.CheckpointPath.GEMMA3_270M_IT + "params": gm.ckpts.CheckpointPath.GEMMA3_270M_IT, }, "gemma3_270m": { "model": gm.nn.Gemma3_270M, - "params": gm.ckpts.CheckpointPath.GEMMA3_270M_PT + "params": gm.ckpts.CheckpointPath.GEMMA3_270M_PT, }, # 1B "gemma3_1b": { From b67cabc39aa256da957ca728e9e8bef641495bfe Mon Sep 17 00:00:00 2001 From: Jyotinder Singh <33001894+JyotinderSingh@users.noreply.github.com> Date: Tue, 16 Sep 2025 12:46:16 +0530 Subject: [PATCH 5/5] removes launch.json --- .vscode/launch.json | 29 ----------------------------- 1 file changed, 29 deletions(-) delete mode 100644 .vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index 8659be2f04..0000000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - // Use IntelliSense to learn about possible attributes. - // Hover to view descriptions of existing attributes. - // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 - "version": "0.2.0", - "configurations": [ - { - "name": "Pytest: Prompt for test target", - "type": "debugpy", - "request": "launch", - "module": "pytest", - "console": "integratedTerminal", - "justMyCode": true, - "args": [ - "${input:pytestTarget}" - ] - } - ], - - // 2. This 'inputs' array defines the prompt - "inputs": [ - { - "id": "pytestTarget", - "type": "promptString", - "description": "Enter the pytest target (file path or file::function)", - "default": "" - } - ] -} \ No newline at end of file