#!/usr/bin/env python3
"""Export the compact release without modifying the FP32 training master.

Usage: python export_fp16.py --source RUN/final --output RUN/final-fp16
Requires the training environment's torch and safetensors packages.
"""

from __future__ import annotations

import argparse
import shutil
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file


def export_fp16(source: Path, output: Path) -> None:
    source = source.resolve()
    output = output.resolve()
    if output == source or source in output.parents:
        raise ValueError("Output must be outside the source model directory")
    if output.exists():
        raise FileExistsError(f"Output already exists: {output}")
    support = ("config_sentence_transformers.json", "modules.json", "tokenizer.json")
    for name in support:
        if not (source / name).is_file():
            raise FileNotFoundError(source / name)
    tensors = load_file(source / "model.safetensors", device="cpu")
    if set(tensors) != {"embedding.weight"}:
        raise ValueError("Expected a single StaticEmbedding weight table")
    weight = tensors["embedding.weight"]
    if weight.ndim != 2 or not weight.is_floating_point():
        raise ValueError("Expected a floating-point embedding matrix")
    converted = weight.to(torch.float16).contiguous()
    if not torch.isfinite(converted).all():
        raise ValueError("FP16 conversion produced non-finite weights")
    output.mkdir(parents=True)
    for name in support:
        shutil.copy2(source / name, output / name)
    if (source / "README.md").is_file():
        shutil.copy2(source / "README.md", output / "README.md")
    save_file({"embedding.weight": converted}, output / "model.safetensors")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    export_fp16(args.source, args.output)
    print(args.output)


if __name__ == "__main__":
    main()