turboderp-org/exllamav2
1
0
Fork 0
mirror of https://github.com/turboderp-org/exllamav2.git synced 2026-03-14 15:57:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix converting files with docker command

Browse Source
This commit is contained in:
Colin
2024-02-20 19:41:57 -05:00
parent c8e2bf4594
commit 008a0bb777
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
util/convert_safetensors.py
View File

@@ -1,15 +1,19 @@
import torch
import argparse, os
import argparse, os, glob
from safetensors.torch import save_file
parser = argparse.ArgumentParser(description="Convert .bin/.pt files to .safetensors")
parser.add_argument("--unshare", action = "store_true", help="Detach tensors to prevent any from sharing memory")
parser.add_argument("input_files", nargs='+', type=str, help="Input file(s)")
parser.add_argument("--unshare", action="store_true", help="Detach tensors to prevent any from sharing memory")
parser.add_argument("input_files", nargs="+", type=str, help="Input file(s)")
args = parser.parse_args()
for file in args.input_files:
tensor_files = []
for file_pattern in args.input_files:
tensor_files.extend(glob.glob(file_pattern))
for file in tensor_files:
print(f" -- Loading {file}...")
state_dict = torch.load(file, map_location = "cpu")
state_dict = torch.load(file, map_location="cpu")
if args.unshare:
for k in state_dict.keys():
@@ -17,4 +21,4 @@ for file in args.input_files:
out_file = os.path.splitext(file)[0] + ".safetensors"
print(f" -- Saving {out_file}...")
save_file(state_dict, out_file, metadata = {'format': 'pt'})
save_file(state_dict, out_file, metadata={"format": "pt"})