| """Convert pinned source weights to strict MLX FP32 bundles. No training or quantization.""" |
| import argparse |
| import importlib.metadata |
| import json |
| import platform |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| sys.path.insert(0, str(ROOT / "scripts")) |
| import numpy as np |
| import mlx.core as mx |
| from safetensors.numpy import load_file |
| from timestral_mlx.config import TimestralConfig |
| from timestral_mlx.model import TimestralModel |
| from download import VARIANTS, sha256 |
|
|
|
|
| def convert(variant): |
| src, dst = ROOT / "models" / "source" / variant, ROOT / "models" / variant |
| source = json.loads((src / "source.json").read_text()) |
| for name, entry in source["files"].items(): |
| if sha256(src / name) != entry["sha256"]: |
| raise ValueError(f"source integrity check failed: {name}") |
| dst.mkdir(parents=True, exist_ok=True) |
| weights = load_file(src / "model.safetensors") |
| if any(v.dtype != np.float32 for v in weights.values()): |
| raise ValueError("expected original FP32 weights") |
| converted = {k.replace("._gate.", ".gate."): mx.array(v) for k, v in weights.items()} |
| config = TimestralConfig.load(src / "config.json") |
| model = TimestralModel(config) |
| model.load_weights(list(converted.items()), strict=True) |
| mx.eval(model.parameters()) |
| mx.save_safetensors(str(dst / "model.safetensors"), converted, |
| metadata={"format": "mlx", "source": source["repo_id"], "revision": source["revision"]}) |
| |
| restored = mx.load(str(dst / "model.safetensors")) |
| if set(restored) != set(converted): |
| raise ValueError("serialized key inventory mismatch") |
| for k, v in weights.items(): |
| if not np.array_equal(v, np.array(restored[k.replace("._gate.", ".gate.")])): |
| raise ValueError(f"lossless conversion check failed: {k}") |
| shutil.copyfile(src / "config.json", dst / "config.json") |
| shutil.copyfile(src / "LICENSE", dst / "LICENSE") |
| manifest = { |
| "format": "timestral-mlx-fp32-v1", "variant": variant, "dtype": "float32", |
| "source": source, "source_code": {"url": "https://github.com/mistralai/timestral", |
| "revision": "a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff"}, |
| "conversion": {"tensor_count": len(weights), "parameter_count": sum(v.size for v in weights.values()), |
| "mapping": "Replace ._gate. with .gate.; tensor layouts and FP32 values unchanged", |
| "lossless_roundtrip_verified": True, "training": False, "quantization": False}, |
| "environment": {"python": platform.python_version(), "mlx": importlib.metadata.version("mlx"), |
| "numpy": np.__version__, "platform": platform.platform()}, |
| "files": {name: {"sha256": sha256(dst / name), "bytes": (dst / name).stat().st_size} |
| for name in ("config.json", "model.safetensors", "LICENSE")}, |
| } |
| (dst / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| print(f"Converted {variant}: {len(weights)} tensors, {manifest['conversion']['parameter_count']:,} parameters; lossless FP32 verified", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("variants", nargs="*", choices=VARIANTS, default=list(VARIANTS)) |
| for size in p.parse_args().variants: |
| convert(size) |
|
|
