File size: 3,461 Bytes
d35e174 | 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 66 67 68 | """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"]})
# Check every serialized tensor, including exact source values and key mapping.
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)
|