Blame for convert_checkpoint.py

"""Convert the owner's checkpoint without enabling unrestricted pickle loading."""
import argparse
from pathlib import Path


def convert(checkpoint, output):
    import torch
    from safetensors.torch import save_file
    import os
    import tempfile
    checkpoint, output = Path(checkpoint), Path(output)
    if output.exists():
        raise ValueError("Output already exists; choose another path.")
    obj = torch.load(checkpoint, map_location="cpu", weights_only=True)
    state = obj.get("state_dict", obj) if isinstance(obj, dict) else None
    if not isinstance(state, dict) or not state:
        raise ValueError("No model tensors were found.")
    if any(not isinstance(k, str) or not isinstance(v, torch.Tensor) for k, v in state.items()):
        raise ValueError("Unexpected checkpoint content. Inspect it before conversion.")
    fd, temp_name = tempfile.mkstemp(dir=output.parent, suffix=".safetensors")
    os.close(fd)
    try:
        save_file({k: v.detach().contiguous().clone() for k, v in state.items()}, temp_name)
        os.replace(temp_name, output)
    finally:
        Path(temp_name).unlink(missing_ok=True)
    return output


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("checkpoint", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    convert(args.checkpoint, args.output)
    print("Saved", args.output)


if __name__ == "__main__":
    main()