chore(clip): update the clip example to make it compatible with HF format

clip/convert.py

@@ -8,6 +8,7 @@
 import mlx.core as mx
 import torch
 from huggingface_hub import snapshot_download
+from mlx_lm.utils import save_weights
 
 
 def get_model_path(path_or_hf_repo: str) -> Path:
@@ -32,44 +33,6 @@ def torch_to_mx(a: torch.Tensor, *, dtype: str) -> mx.array:
     return mx.array(a.numpy(), getattr(mx, dtype))
 
 
-def map_weights(key: str, value: torch.Tensor) -> Tuple[str, mx.array]:
-    key = key.replace("embeddings.", "")
-    key = key.replace("encoder.", "")
-    key = key.replace("position_embedding.weight", "position_embedding")
-
-    # Map attention layers
-    if "self_attn." in key:
-        key = key.replace("self_attn.", "attention.")
-    if "q_proj." in key:
-        key = key.replace("q_proj.", "query_proj.")
-    if "k_proj." in key:
-        key = key.replace("k_proj.", "key_proj.")
-    if "v_proj." in key:
-        key = key.replace("v_proj.", "value_proj.")
-    if "layer_norm1." in key:
-        key = key.replace("layer_norm1.", "ln1.")
-    if "layer_norm2." in key:
-        key = key.replace("layer_norm2.", "ln2.")
-    # Map ffn layers
-    if "mlp.fc1" in key:
-        key = key.replace("mlp.fc1", "linear1")
-    if "mlp.fc2" in key:
-        key = key.replace("mlp.fc2", "linear2")
-    # Fix layernorm typo
-    if "pre_layrnorm" in key:
-        # Fix typo in weights :)
-        key = key.replace("pre_layrnorm", "pre_layernorm")
-    if "patch_embedding.weight" in key:
-        # Initially, value: [out_channels, in_channels, kH, KW].
-        # We want [out_channels, kH, KW, in_channels]
-        value = value.permute(0, 2, 3, 1)
-    return (key, torch_to_mx(value, dtype=str(value.dtype).replace("torch.", "")))
-
-
-def should_keep_weight(key: str):
-    return not ("position_ids" in key)
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Download and Convert (OpenAI) CLIP weights to MLX"
@@ -86,7 +49,13 @@ def should_keep_weight(key: str):
         default="mlx_model",
         help="Path to save the MLX model.",
     )
-
+    parser.add_argument(
+        "--dtype",
+        help="dtype for loading the torch model and input for quantization or saving the converted model. "
+        "The original weights are stored in bfloat16.",
+        type=str,
+        default="float32",
+    )
     args = parser.parse_args()
 
     torch_path = get_model_path(args.hf_repo)
@@ -96,10 +65,11 @@ def should_keep_weight(key: str):
     print("[INFO] Loading")
     torch_weights = torch.load(torch_path / "pytorch_model.bin")
     print("[INFO] Converting")
-    mlx_weights = dict(map_weights(k, v) for (k, v) in torch_weights.items())
-    mlx_weights = {k: v for (k, v) in mlx_weights.items() if should_keep_weight(k)}
+    mlx_weights = {
+        k: torch_to_mx(v, dtype=args.dtype) for k, v in torch_weights.items()
+    }
     print("[INFO] Saving")
-    mx.savez(str(mlx_path / "weights.npz"), **mlx_weights)
+    save_weights(mlx_path, mlx_weights)
     for fn in ["config.json", "merges.txt", "vocab.json", "preprocessor_config.json"]:
         shutil.copyfile(
             str(torch_path / f"{fn}"),