update script

convert_flax.py

@@ -127,20 +127,8 @@ def convert_to_hf(path: Path):
     print(f"{num_layers=}")
     print(f"{num_siglip_layers=}")
 
-    def load_params(*keys: tuple[str, ...], prefix: str | None = None):
-        # load params with specific keys and params starts with prefix
-        f1 = lambda k: tuple(subkey.key for subkey in k) in keys
-        f2 = lambda k: k[0].key.startswith(prefix)
-
-        # set to None to not load that weights
-        pytree = jax.tree.map_with_path(lambda k, v: v if f1(k) or f2(k) else None, metadata)
-        return ckpt.restore(path, pytree)
-
     # NOTE: all gemma3 models use tied embeddings, even for the 27B version.
-    params = load_params(
-        ("transformer/final_norm", "scale"),
-        prefix="transformer/embedder",
-    )
+    params = ckpt.restore(path)
     state_dict = dict()
 
     if num_siglip_layers > 0:
@@ -164,7 +152,6 @@ def convert_to_hf(path: Path):
 
     for layer_idx in range(num_layers):
         jax_prefix = f"transformer/layer_{layer_idx}/"
-        params = load_params(prefix=jax_prefix)
 
         state_dict = dict()
         prefix = f"{gemma_prefix}model.layers.{layer_idx}."
@@ -200,7 +187,6 @@ def convert_to_hf(path: Path):
 
     # vision tower
     if num_siglip_layers > 0:
-        params = load_params(prefix=SIGLIP_PREFIX)
         siglip_state_dict = convert_siglip(params, num_siglip_layers)
         for k, v in siglip_state_dict.items():
             state_dict[f"vision_tower.vision_model.{k}"] = v
@@ -272,21 +258,22 @@ if __name__ == "__main__":
     filename = f"model-{shard_idx + 1:05d}.safetensors"
     for sub_state_dict in tqdm(convert_to_hf(args.ckpt_dir)):
         sub_state_dict = convert_awq(sub_state_dict)
+        new_size = sum(v.nbytes for v in sub_state_dict.values())
 
-        for k, v in sub_state_dict.items():
-            state_dict[k] = v
-            size += v.nbytes
-
-            total_size += v.nbytes
-            weight_map[k] = filename
-
-        if size > 5e9:
+        if size + new_size > 5e9:
             save_file(state_dict, args.save_dir / filename)
             state_dict = dict()
             size = 0
             shard_idx += 1
             filename = f"model-{shard_idx + 1:05d}.safetensors"
 
+        # assume that new_size < 5e9
+        size += new_size
+        total_size += new_size
+        for k, v in sub_state_dict.items():
+            state_dict[k] = v
+            weight_map[k] = filename
+
     save_file(state_dict, args.save_dir / filename)
     json.dump(
         dict(metadata=dict(total_size=total_size), weight_map=weight_map),