usrbinenv python 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert a pytorchmodel bin or model pt file to torch float16 for faster downloads less disk space statedict torch loadsrcpath maplocationmaplocation for k v in tqdmstatedict items if not isinstancev torch tensor raise typeerrorfp16 conversion only works on paths that are saved state dicts like pytorchmodel bin statedictk v half if savepath is none overwrite srcpath savepath srcpath torch savestatedict savepath if name main fire fireconvert usr bin env python 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert a pytorch_model bin or model pt file to torch float16 for faster downloads less disk space overwrite src_path
from typing import Union

import fire
import torch
from tqdm import tqdm


def convert(src_path: str, map_location: str = "cpu", save_path: Union[str, None] = None) -> None:
    
    state_dict = torch.load(src_path, map_location=map_location)
    for k, v in tqdm(state_dict.items()):
        if not isinstance(v, torch.Tensor):
            raise TypeError("FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin")
        state_dict[k] = v.half()
    if save_path is None:  
        save_path = src_path
    torch.save(state_dict, save_path)


if __name__ == "__main__":
    fire.Fire(convert)