From 0593a2b9a5d9232e4ae6e666ab44b559838ab97d Mon Sep 17 00:00:00 2001 From: Mansi Mehta Date: Fri, 17 Oct 2025 06:06:12 +0530 Subject: [PATCH] Update preset map values for paligemma2 --- .../convert_pali_gemma2_checkpoints.py | 38 +++++++++---------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/tools/checkpoint_conversion/convert_pali_gemma2_checkpoints.py b/tools/checkpoint_conversion/convert_pali_gemma2_checkpoints.py index 23bda97b72..e99ad80cde 100644 --- a/tools/checkpoint_conversion/convert_pali_gemma2_checkpoints.py +++ b/tools/checkpoint_conversion/convert_pali_gemma2_checkpoints.py @@ -63,27 +63,23 @@ FLAGS = flags.FLAGS PRESET_MAP = { - "pali_gemma2_3b_ft_docci_448": ( - "google/paligemma-2/jax/paligemma2-3b-ft-docci-448" - ), - "pali_gemma2_10b_ft_docci_448": ( - "google/paligemma-2/jax/paligemma2-10b-ft-docci-448" - ), - "pali_gemma2_3b_mix_224": "google/paligemma-2/jax/paligemma2-3b-mix-224", - "pali_gemma2_3b_mix_448": "google/paligemma-2/jax/paligemma2-3b-mix-448", - "pali_gemma2_10b_mix_224": "google/paligemma-2/jax/paligemma2-10b-mix-224", - "pali_gemma2_10b_mix_448": "google/paligemma-2/jax/paligemma2-10b-mix-448", - "pali_gemma2_28b_mix_224": "google/paligemma-2/jax/paligemma2-28b-mix-224", - "pali_gemma2_28b_mix_448": "google/paligemma-2/jax/paligemma2-28b-mix-448", - "pali_gemma2_3b_pt_224": "google/paligemma-2/jax/paligemma2-3b-pt-224", - "pali_gemma2_3b_pt_448": "google/paligemma-2/jax/paligemma2-3b-pt-448", - "pali_gemma2_3b_pt_896": "google/paligemma-2/jax/paligemma2-3b-pt-896", - "pali_gemma2_10b_pt_224": "google/paligemma-2/jax/paligemma2-10b-pt-224", - "pali_gemma2_10b_pt_448": "google/paligemma-2/jax/paligemma2-10b-pt-448", - "pali_gemma2_10b_pt_896": "google/paligemma-2/jax/paligemma2-10b-pt-896", - "pali_gemma2_28b_pt_224": "google/paligemma-2/jax/paligemma2-28b-pt-224", - "pali_gemma2_28b_pt_448": "google/paligemma-2/jax/paligemma2-28b-pt-448", - "pali_gemma2_28b_pt_896": "google/paligemma-2/jax/paligemma2-28b-pt-896", + "pali_gemma_2_ft_docci_3b_448": "google/paligemma2-3b-ft-docci-448-jax", + "pali_gemma2_ft_docci_10b_448": "google/paligemma2-10b-ft-docci-448", + "pali_gemma2_mix_3b_224": "google/paligemma2-3b-mix-224", + "pali_gemma2_mix_3b_448": "google/paligemma2-3b-mix-448", + "pali_gemma2_mix_10b_224": "google/paligemma2-10b-mix-224", + "pali_gemma2_mix_10b_448": "google/paligemma2-10b-mix-448-jax", + "pali_gemma2_mix_28b_224": "google/paligemma2-28b-mix-224", + "pali_gemma2_mix_28b_448": "google/paligemma2-28b-mix-448-jax", + "pali_gemma2_pt_3b_224": "google/paligemma2-3b-pt-224-jax", + "pali_gemma2_pt_3b_448": "google/paligemma2-3b-pt-448", + "pali_gemma2_pt_3b_896": "google/paligemma2-3b-pt-896", + "pali_gemma2_pt_10b_224": "google/paligemma2-10b-pt-224", + "pali_gemma2_pt_10b_448": "google/paligemma2-10b-pt-448", + "pali_gemma2_pt_10b_896": "google/paligemma2-10b-pt-896", + "pali_gemma2_pt_28b_224": "google/paligemma2-28b-pt-224", + "pali_gemma2_pt_28b_448": "google/paligemma2-28b-pt-448", + "pali_gemma2_pt_28b_896": "google/paligemma2-28b-pt-896", }