Upload export_apochat_litert.py with huggingface_hub

export_apochat_litert.py

@@ -134,6 +134,21 @@ def download_repo_files(repo_id: str, revision: str | None, local_dir: Path) ->
         )
 
 
+def mlx_key_to_pytorch_key(key: str) -> str:
+    """Map MLX Gemma 4 key names to PyTorch / transformers key names."""
+    if key.startswith("language_model.model."):
+        return "model.language_model." + key[len("language_model.model."):]
+    if key.startswith("audio_tower."):
+        return "model.audio_tower." + key[len("audio_tower."):]
+    if key.startswith("vision_tower."):
+        return "model.vision_tower." + key[len("vision_tower."):]
+    if key.startswith("embed_audio.embedding_projection"):
+        return key.replace("embed_audio.embedding_projection", "model.embed_audio.embedding_projection", 1)
+    if key.startswith("embed_vision.embedding_projection"):
+        return key.replace("embed_vision.embedding_projection", "model.embed_vision.embedding_projection", 1)
+    raise ValueError(f"Unexpected MLX key prefix: {key}")
+
+
 def patch_config_for_pytorch(config_path: Path) -> None:
     """Remove MLX quantization config and ensure torch_dtype is bfloat16."""
     with open(config_path, "r", encoding="utf-8") as f:
@@ -218,7 +233,8 @@ def dequantize_mlx_to_pytorch(
             arr = arr.astype(mx.bfloat16)
 
         torch_tensor = mlx_bfloat16_to_torch(arr)
-        current_shard[name] = torch_tensor
+        pytorch_name = mlx_key_to_pytorch_key(name)
+        current_shard[pytorch_name] = torch_tensor
         current_shard_bytes += torch_tensor.nbytes
 
         if current_shard_bytes >= shard_size_bytes: