| """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() |
|
|
