| """ | |
| Usage: | |
| python3 -m fastchat.model.convert_fp16 --in in-folder --out out-folder | |
| """ | |
| import argparse | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import torch | |
| def convert_fp16(in_checkpoint, out_checkpoint): | |
| tokenizer = AutoTokenizer.from_pretrained(in_checkpoint, use_fast=False) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| in_checkpoint, torch_dtype=torch.float16, low_cpu_mem_usage=True | |
| ) | |
| model.save_pretrained(out_checkpoint) | |
| tokenizer.save_pretrained(out_checkpoint) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--in-checkpoint", type=str, help="Path to the model") | |
| parser.add_argument("--out-checkpoint", type=str, help="Path to the output model") | |
| args = parser.parse_args() | |
| convert_fp16(args.in_checkpoint, args.out_checkpoint) | |