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