File size: 2,050 Bytes
591cda4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """Losslessly convert the upstream FP32 checkpoint to native MLX layout."""
from __future__ import annotations
import argparse
import json
import shutil
from pathlib import Path
import mlx.core as mx
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("source", type=Path)
parser.add_argument("output", type=Path)
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=True)
weights = mx.load(str(args.source / "model.safetensors"))
converted = {}
for name, tensor in weights.items():
if name.endswith("conv.conv.weight"):
if tensor.ndim != 3 or tensor.shape[1:] != (1, 3):
raise ValueError(f"Unexpected convolution shape for {name}: {tensor.shape}")
tensor = tensor.transpose(0, 2, 1)
if tensor.dtype != mx.float32:
raise TypeError(f"Expected FP32 source tensor {name}, got {tensor.dtype}")
converted[name] = tensor
mx.save_safetensors(
str(args.output / "model.safetensors"),
converted,
metadata={
"format": "mlx",
"source": "LiquidAI/LFM2.5-Encoder-350M",
"source_revision": "b886781f7c6f10ca9b7096e21b83e30a073c2f39",
"dtype": "float32",
},
)
for filename in (
"config.json",
"tokenizer.json",
"tokenizer_config.json",
"special_tokens_map.json",
"LICENSE",
):
shutil.copy2(args.source / filename, args.output / filename)
manifest = {
"source": "LiquidAI/LFM2.5-Encoder-350M",
"source_revision": "b886781f7c6f10ca9b7096e21b83e30a073c2f39",
"parameter_dtype": "float32",
"parameter_count": int(sum(tensor.size for tensor in converted.values())),
"tensor_count": len(converted),
"transforms": ["conv weights [out,1,kernel] -> [out,kernel,1]"],
}
(args.output / "conversion_manifest.json").write_text(
json.dumps(manifest, indent=2) + "\n"
)
if __name__ == "__main__":
main()
|