Add znn
Browse files- .gitattributes +1 -0
- model-00001-of-00005.safetensors → model-00001-of-00005.safetensors.znn +2 -2
- model-00002-of-00005.safetensors → model-00002-of-00005.safetensors.znn +2 -2
- model-00003-of-00005.safetensors → model-00003-of-00005.safetensors.znn +2 -2
- model-00004-of-00005.safetensors → model-00004-of-00005.safetensors.znn +2 -2
- model-00005-of-00005.safetensors +0 -3
- model-00005-of-00005.safetensors.znn +3 -0
- model.safetensors.index.json +0 -0
- zipnn_compress_file.py +179 -0
- zipnn_compress_path.py +380 -0
- zipnn_decompress_file.py +101 -0
- zipnn_decompress_path.py +303 -0
.gitattributes
CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
36 |
+
*.znn filter=lfs diff=lfs merge=lfs -text
|
model-00001-of-00005.safetensors → model-00001-of-00005.safetensors.znn
RENAMED
@@ -1,3 +1,3 @@
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:
|
3 |
-
size
|
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:35fdc0739513b5196075e8b687078e0e2ff25a78c355316aaf482e460887ca4a
|
3 |
+
size 3282411963
|
model-00002-of-00005.safetensors → model-00002-of-00005.safetensors.znn
RENAMED
@@ -1,3 +1,3 @@
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:
|
3 |
-
size
|
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:f44c7f7e7bc454b029dcd79bea70619a16795ea93e0c3dad4536c72fbcd16617
|
3 |
+
size 3296284300
|
model-00003-of-00005.safetensors → model-00003-of-00005.safetensors.znn
RENAMED
@@ -1,3 +1,3 @@
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:
|
3 |
-
size
|
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:5ead131bb7dcbb3ab271bb4527fb5deb0c06e919261db6f6e0c21765882a1402
|
3 |
+
size 3260738442
|
model-00004-of-00005.safetensors → model-00004-of-00005.safetensors.znn
RENAMED
@@ -1,3 +1,3 @@
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:
|
3 |
-
size
|
|
|
1 |
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:ada192d81d1a3971af1febf54601d07dc65a1702684d1f2eddd3236256df3a1d
|
3 |
+
size 3313808717
|
model-00005-of-00005.safetensors
DELETED
@@ -1,3 +0,0 @@
|
|
1 |
-
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:656566a9aeee1115ca174c16ea3c51a2839b3a7f14cc349805aa0ef5c7433e7b
|
3 |
-
size 1465943128
|
|
|
|
|
|
|
|
model-00005-of-00005.safetensors.znn
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:d23e328525f5fc52d04f55a16f034e515373b1fcc7ca6c0fa2954ac0b6d10f8e
|
3 |
+
size 974361964
|
model.safetensors.index.json
CHANGED
The diff for this file is too large to render.
See raw diff
|
|
zipnn_compress_file.py
ADDED
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import subprocess
|
3 |
+
import sys
|
4 |
+
import argparse
|
5 |
+
import time
|
6 |
+
|
7 |
+
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
8 |
+
|
9 |
+
KB = 1024
|
10 |
+
MB = 1024 * 1024
|
11 |
+
GB = 1024 * 1024 * 1024
|
12 |
+
|
13 |
+
RED = "\033[91m"
|
14 |
+
YELLOW = "\033[93m"
|
15 |
+
GREEN = "\033[92m"
|
16 |
+
RESET = "\033[0m"
|
17 |
+
|
18 |
+
def check_and_install_zipnn():
|
19 |
+
try:
|
20 |
+
import zipnn
|
21 |
+
except ImportError:
|
22 |
+
print("zipnn not found. Installing...")
|
23 |
+
subprocess.check_call(
|
24 |
+
[
|
25 |
+
sys.executable,
|
26 |
+
"-m",
|
27 |
+
"pip",
|
28 |
+
"install",
|
29 |
+
"zipnn",
|
30 |
+
"--upgrade",
|
31 |
+
]
|
32 |
+
)
|
33 |
+
import zipnn
|
34 |
+
|
35 |
+
|
36 |
+
def parse_streaming_chunk_size(
|
37 |
+
streaming_chunk_size,
|
38 |
+
):
|
39 |
+
if str(streaming_chunk_size).isdigit():
|
40 |
+
final = int(streaming_chunk_size)
|
41 |
+
else:
|
42 |
+
size_value = int(streaming_chunk_size[:-2])
|
43 |
+
size_unit = streaming_chunk_size[-2].lower()
|
44 |
+
|
45 |
+
if size_unit == "k":
|
46 |
+
final = KB * size_value
|
47 |
+
elif size_unit == "m":
|
48 |
+
final = MB * size_value
|
49 |
+
elif size_unit == "g":
|
50 |
+
final = GB * size_value
|
51 |
+
else:
|
52 |
+
raise ValueError(f"Invalid size unit: {size_unit}. Use 'k', 'm', or 'g'.")
|
53 |
+
|
54 |
+
return final
|
55 |
+
|
56 |
+
|
57 |
+
def compress_file(
|
58 |
+
input_file,
|
59 |
+
dtype="",
|
60 |
+
streaming_chunk_size=1048576,
|
61 |
+
delete=False,
|
62 |
+
force=False,
|
63 |
+
hf_cache=False,
|
64 |
+
):
|
65 |
+
import zipnn
|
66 |
+
|
67 |
+
streaming_chunk_size = parse_streaming_chunk_size(streaming_chunk_size)
|
68 |
+
full_path = input_file
|
69 |
+
if not os.path.exists(full_path):
|
70 |
+
print(f"{RED}File not found{RESET}")
|
71 |
+
return
|
72 |
+
|
73 |
+
compressed_path = full_path + ".znn"
|
74 |
+
if not force and os.path.exists(compressed_path):
|
75 |
+
user_input = (
|
76 |
+
input(f"{compressed_path} already exists; overwrite (y/n)? ").strip().lower()
|
77 |
+
)
|
78 |
+
if user_input not in ("yes", "y"):
|
79 |
+
print(f"Skipping {full_path}...")
|
80 |
+
return
|
81 |
+
print(f"Compressing {full_path}...")
|
82 |
+
#
|
83 |
+
output_file = input_file + ".znn"
|
84 |
+
if dtype:
|
85 |
+
zpn = zipnn.ZipNN(
|
86 |
+
bytearray_dtype="float32",
|
87 |
+
is_streaming=True,
|
88 |
+
streaming_chunk_kb=streaming_chunk_size,
|
89 |
+
)
|
90 |
+
else:
|
91 |
+
zpn = zipnn.ZipNN(
|
92 |
+
is_streaming=True,
|
93 |
+
streaming_chunk_kb=streaming_chunk_size,
|
94 |
+
)
|
95 |
+
file_size_before = 0
|
96 |
+
file_size_after = 0
|
97 |
+
start_time = time.time()
|
98 |
+
with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
|
99 |
+
chunk = infile.read()
|
100 |
+
file_size_before += len(chunk)
|
101 |
+
compressed_chunk = zpn.compress(chunk)
|
102 |
+
if compressed_chunk:
|
103 |
+
file_size_after += len(compressed_chunk)
|
104 |
+
outfile.write(compressed_chunk)
|
105 |
+
end_time = time.time() - start_time
|
106 |
+
print(f"Compressed {input_file} to {output_file}")
|
107 |
+
print(
|
108 |
+
f"{GREEN}Original size: {file_size_before/GB:.02f}GB size after compression: {file_size_after/GB:.02f}GB, Remaining size is {file_size_after/file_size_before*100:.02f}% of original, time: {end_time:.02f}{RESET}"
|
109 |
+
)
|
110 |
+
|
111 |
+
if delete and not hf_cache:
|
112 |
+
print(f"Deleting {full_path}...")
|
113 |
+
os.remove(full_path)
|
114 |
+
|
115 |
+
if hf_cache:
|
116 |
+
# If the file is in the Hugging Face cache, fix the symlinks
|
117 |
+
print(f"{YELLOW}Reorganizing Hugging Face cache...{RESET}")
|
118 |
+
try:
|
119 |
+
snapshot_path = os.path.dirname(input_file)
|
120 |
+
blob_name = os.path.join(snapshot_path, os.readlink(input_file))
|
121 |
+
os.rename(output_file, blob_name)
|
122 |
+
os.symlink(blob_name, output_file)
|
123 |
+
if os.path.exists(input_file):
|
124 |
+
os.remove(input_file)
|
125 |
+
except Exception as e:
|
126 |
+
raise Exception(f"Error reorganizing Hugging Face cache: {e}")
|
127 |
+
|
128 |
+
if __name__ == "__main__":
|
129 |
+
if len(sys.argv) < 2:
|
130 |
+
print("Usage: python compress_files.py <suffix>")
|
131 |
+
print("Example: python compress_files.py 'safetensors'")
|
132 |
+
sys.exit(1)
|
133 |
+
|
134 |
+
parser = argparse.ArgumentParser(description="Enter a file path to compress.")
|
135 |
+
parser.add_argument(
|
136 |
+
"input_file",
|
137 |
+
type=str,
|
138 |
+
help="Specify the path to the file to compress.",
|
139 |
+
)
|
140 |
+
parser.add_argument(
|
141 |
+
"--float32",
|
142 |
+
action="store_true",
|
143 |
+
help="A flag that triggers float32 compression",
|
144 |
+
)
|
145 |
+
parser.add_argument(
|
146 |
+
"--streaming_chunk_size",
|
147 |
+
type=str,
|
148 |
+
help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
|
149 |
+
)
|
150 |
+
parser.add_argument(
|
151 |
+
"--delete",
|
152 |
+
action="store_true",
|
153 |
+
help="A flag that triggers deletion of a single file instead of compression",
|
154 |
+
)
|
155 |
+
parser.add_argument(
|
156 |
+
"--force",
|
157 |
+
action="store_true",
|
158 |
+
help="A flag that forces overwriting when compressing.",
|
159 |
+
)
|
160 |
+
parser.add_argument(
|
161 |
+
"--hf_cache",
|
162 |
+
action="store_true",
|
163 |
+
help="A flag that indicates if the file is in the Hugging Face cache.",
|
164 |
+
)
|
165 |
+
args = parser.parse_args()
|
166 |
+
optional_kwargs = {}
|
167 |
+
if args.float32:
|
168 |
+
optional_kwargs["dtype"] = 32
|
169 |
+
if args.streaming_chunk_size is not None:
|
170 |
+
optional_kwargs["streaming_chunk_size"] = args.streaming_chunk_size
|
171 |
+
if args.delete:
|
172 |
+
optional_kwargs["delete"] = args.delete
|
173 |
+
if args.force:
|
174 |
+
optional_kwargs["force"] = args.force
|
175 |
+
if args.hf_cache:
|
176 |
+
optional_kwargs["hf_cache"] = args.hf_cache
|
177 |
+
|
178 |
+
check_and_install_zipnn()
|
179 |
+
compress_file(args.input_file, **optional_kwargs)
|
zipnn_compress_path.py
ADDED
@@ -0,0 +1,380 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import subprocess
|
3 |
+
import sys
|
4 |
+
import argparse
|
5 |
+
from pathlib import Path
|
6 |
+
from concurrent.futures import (
|
7 |
+
ProcessPoolExecutor,
|
8 |
+
as_completed,
|
9 |
+
)
|
10 |
+
from zipnn_compress_file import compress_file
|
11 |
+
|
12 |
+
sys.path.append(
|
13 |
+
os.path.abspath(
|
14 |
+
os.path.join(
|
15 |
+
os.path.dirname(__file__), ".."
|
16 |
+
)
|
17 |
+
)
|
18 |
+
)
|
19 |
+
|
20 |
+
|
21 |
+
KB = 1024
|
22 |
+
MB = 1024 * 1024
|
23 |
+
GB = 1024 * 1024 * 1024
|
24 |
+
|
25 |
+
RED = "\033[91m"
|
26 |
+
YELLOW = "\033[93m"
|
27 |
+
GREEN = "\033[92m"
|
28 |
+
RESET = "\033[0m"
|
29 |
+
|
30 |
+
|
31 |
+
def check_and_install_zipnn():
|
32 |
+
try:
|
33 |
+
import zipnn
|
34 |
+
except ImportError:
|
35 |
+
print("zipnn not found. Installing...")
|
36 |
+
subprocess.check_call(
|
37 |
+
[
|
38 |
+
sys.executable,
|
39 |
+
"-m",
|
40 |
+
"pip",
|
41 |
+
"install",
|
42 |
+
"zipnn",
|
43 |
+
"--upgrade",
|
44 |
+
]
|
45 |
+
)
|
46 |
+
import zipnn
|
47 |
+
|
48 |
+
|
49 |
+
def parse_streaming_chunk_size(
|
50 |
+
streaming_chunk_size,
|
51 |
+
):
|
52 |
+
if str(streaming_chunk_size).isdigit():
|
53 |
+
final = int(streaming_chunk_size)
|
54 |
+
else:
|
55 |
+
size_value = int(
|
56 |
+
streaming_chunk_size[:-2]
|
57 |
+
)
|
58 |
+
size_unit = streaming_chunk_size[
|
59 |
+
-2
|
60 |
+
].lower()
|
61 |
+
|
62 |
+
if size_unit == "k":
|
63 |
+
final = KB * size_value
|
64 |
+
elif size_unit == "m":
|
65 |
+
final = MB * size_value
|
66 |
+
elif size_unit == "g":
|
67 |
+
final = GB * size_value
|
68 |
+
else:
|
69 |
+
raise ValueError(
|
70 |
+
f"Invalid size unit: {size_unit}. Use 'k', 'm', or 'g'."
|
71 |
+
)
|
72 |
+
|
73 |
+
return final
|
74 |
+
|
75 |
+
def replace_in_file(file_path, old: str, new: str) -> None:
|
76 |
+
"""Given a file_path, replace all occurrences of `old` with `new` inpalce."""
|
77 |
+
|
78 |
+
with open(file_path, 'r') as file:
|
79 |
+
file_data = file.read()
|
80 |
+
|
81 |
+
file_data = file_data.replace(old, new)
|
82 |
+
|
83 |
+
with open(file_path, 'w') as file:
|
84 |
+
file.write(file_data)
|
85 |
+
|
86 |
+
def compress_files_with_suffix(
|
87 |
+
suffix,
|
88 |
+
dtype="",
|
89 |
+
streaming_chunk_size=1048576,
|
90 |
+
path=".",
|
91 |
+
delete=False,
|
92 |
+
r=False,
|
93 |
+
force=False,
|
94 |
+
max_processes=1,
|
95 |
+
hf_cache=False,
|
96 |
+
model="",
|
97 |
+
branch="main",
|
98 |
+
):
|
99 |
+
import zipnn
|
100 |
+
|
101 |
+
overwrite_first=True
|
102 |
+
file_list = []
|
103 |
+
streaming_chunk_size = (
|
104 |
+
parse_streaming_chunk_size(
|
105 |
+
streaming_chunk_size
|
106 |
+
)
|
107 |
+
)
|
108 |
+
if model:
|
109 |
+
if not hf_cache:
|
110 |
+
raise ValueError(
|
111 |
+
"Must specify --hf_cache when using --model"
|
112 |
+
)
|
113 |
+
try:
|
114 |
+
from huggingface_hub import scan_cache_dir
|
115 |
+
except ImportError:
|
116 |
+
raise ImportError(
|
117 |
+
"huggingface_hub not found. Please pip install huggingface_hub."
|
118 |
+
)
|
119 |
+
cache = scan_cache_dir()
|
120 |
+
repo = next((repo for repo in cache.repos if repo.repo_id == model), None)
|
121 |
+
|
122 |
+
if repo is not None:
|
123 |
+
print(f"Found repo {model} in cache")
|
124 |
+
|
125 |
+
# Get the latest revision path
|
126 |
+
hash = ''
|
127 |
+
try:
|
128 |
+
with open(os.path.join(repo.repo_path, 'refs', branch), "r") as ref:
|
129 |
+
hash = ref.read()
|
130 |
+
except FileNotFoundError:
|
131 |
+
raise FileNotFoundError(f"Branch {branch} not found in repo {model}")
|
132 |
+
|
133 |
+
path = os.path.join(repo.repo_path, 'snapshots', hash)
|
134 |
+
|
135 |
+
directories_to_search = (
|
136 |
+
os.walk(path)
|
137 |
+
if r
|
138 |
+
else [(path, [], os.listdir(path))]
|
139 |
+
)
|
140 |
+
files_found = False
|
141 |
+
for root, _, files in directories_to_search:
|
142 |
+
for file_name in files:
|
143 |
+
if file_name.endswith(suffix):
|
144 |
+
compressed_path = (
|
145 |
+
file_name + ".znn"
|
146 |
+
)
|
147 |
+
if not force and os.path.exists(
|
148 |
+
compressed_path
|
149 |
+
):
|
150 |
+
#
|
151 |
+
if overwrite_first:
|
152 |
+
overwrite_first=False
|
153 |
+
user_input = (
|
154 |
+
input(
|
155 |
+
f"Compressed files already exists; Would you like to overwrite them all (y/n)? "
|
156 |
+
)
|
157 |
+
.strip()
|
158 |
+
.lower()
|
159 |
+
)
|
160 |
+
if user_input not in (
|
161 |
+
"y",
|
162 |
+
"yes",
|
163 |
+
):
|
164 |
+
print(
|
165 |
+
f"No forced overwriting."
|
166 |
+
)
|
167 |
+
else:
|
168 |
+
print(
|
169 |
+
f"Overwriting all compressed files."
|
170 |
+
)
|
171 |
+
force=True
|
172 |
+
#
|
173 |
+
if not force:
|
174 |
+
user_input = (
|
175 |
+
input(
|
176 |
+
f"{compressed_path} already exists; overwrite (y/n)? "
|
177 |
+
)
|
178 |
+
.strip()
|
179 |
+
.lower()
|
180 |
+
)
|
181 |
+
if user_input not in (
|
182 |
+
"y",
|
183 |
+
"yes",
|
184 |
+
):
|
185 |
+
print(
|
186 |
+
f"Skipping {file_name}..."
|
187 |
+
)
|
188 |
+
continue
|
189 |
+
files_found = True
|
190 |
+
full_path = os.path.join(
|
191 |
+
root, file_name
|
192 |
+
)
|
193 |
+
file_list.append(full_path)
|
194 |
+
|
195 |
+
if file_list and hf_cache:
|
196 |
+
try:
|
197 |
+
from transformers.utils import (
|
198 |
+
SAFE_WEIGHTS_INDEX_NAME,
|
199 |
+
WEIGHTS_INDEX_NAME
|
200 |
+
)
|
201 |
+
except ImportError:
|
202 |
+
raise ImportError(
|
203 |
+
"Transformers not found. Please pip install transformers."
|
204 |
+
)
|
205 |
+
|
206 |
+
if os.path.exists(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)):
|
207 |
+
print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
|
208 |
+
blob_name = os.path.join(path, os.readlink(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)))
|
209 |
+
replace_in_file(
|
210 |
+
file_path=blob_name,
|
211 |
+
old=f"{suffix}",
|
212 |
+
new=f"{suffix}.znn"
|
213 |
+
)
|
214 |
+
elif os.path.exists(os.path.join(path, WEIGHTS_INDEX_NAME)):
|
215 |
+
print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
|
216 |
+
blob_name = os.path.join(path, os.readlink(os.path.join(path, WEIGHTS_INDEX_NAME)))
|
217 |
+
replace_in_file(
|
218 |
+
file_path=blob_name,
|
219 |
+
old=f"{suffix}",
|
220 |
+
new=f"{suffix}.znn"
|
221 |
+
)
|
222 |
+
|
223 |
+
with ProcessPoolExecutor(
|
224 |
+
max_workers=max_processes
|
225 |
+
) as executor:
|
226 |
+
future_to_file = {
|
227 |
+
executor.submit(
|
228 |
+
compress_file,
|
229 |
+
file,
|
230 |
+
dtype,
|
231 |
+
streaming_chunk_size,
|
232 |
+
delete,
|
233 |
+
True,
|
234 |
+
hf_cache,
|
235 |
+
): file
|
236 |
+
for file in file_list[:max_processes]
|
237 |
+
}
|
238 |
+
file_list = file_list[max_processes:]
|
239 |
+
while future_to_file:
|
240 |
+
for future in as_completed(
|
241 |
+
future_to_file
|
242 |
+
):
|
243 |
+
file = future_to_file.pop(future)
|
244 |
+
|
245 |
+
try:
|
246 |
+
future.result()
|
247 |
+
except Exception as exc:
|
248 |
+
print(
|
249 |
+
f"{RED}File {file} generated an exception: {exc}{RESET}"
|
250 |
+
)
|
251 |
+
|
252 |
+
if file_list:
|
253 |
+
next_file = file_list.pop(0)
|
254 |
+
future_to_file[
|
255 |
+
executor.submit(
|
256 |
+
compress_file,
|
257 |
+
next_file,
|
258 |
+
dtype,
|
259 |
+
streaming_chunk_size,
|
260 |
+
delete,
|
261 |
+
True,
|
262 |
+
hf_cache,
|
263 |
+
)
|
264 |
+
] = next_file
|
265 |
+
|
266 |
+
if not files_found:
|
267 |
+
print(
|
268 |
+
f"{RED}No files with the suffix '{suffix}' found.{RESET}"
|
269 |
+
)
|
270 |
+
|
271 |
+
print(f"{GREEN}All files compressed{RESET}")
|
272 |
+
|
273 |
+
|
274 |
+
if __name__ == "__main__":
|
275 |
+
if len(sys.argv) < 2:
|
276 |
+
print(
|
277 |
+
"Usage: python compress_files.py <suffix>"
|
278 |
+
)
|
279 |
+
print(
|
280 |
+
"Example: python compress_files.py 'safetensors'"
|
281 |
+
)
|
282 |
+
sys.exit(1)
|
283 |
+
|
284 |
+
parser = argparse.ArgumentParser(
|
285 |
+
description="Enter a suffix to compress, (optional) dtype, (optional) streaming chunk size, (optional) path to files."
|
286 |
+
)
|
287 |
+
parser.add_argument(
|
288 |
+
"suffix",
|
289 |
+
type=str,
|
290 |
+
help="Specify the file suffix to compress all files with that suffix. If a single file name is provided, only that file will be compressed.",
|
291 |
+
)
|
292 |
+
parser.add_argument(
|
293 |
+
"--float32",
|
294 |
+
action="store_true",
|
295 |
+
help="A flag that triggers float32 compression",
|
296 |
+
)
|
297 |
+
parser.add_argument(
|
298 |
+
"--streaming_chunk_size",
|
299 |
+
type=str,
|
300 |
+
help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
|
301 |
+
)
|
302 |
+
parser.add_argument(
|
303 |
+
"--path",
|
304 |
+
type=str,
|
305 |
+
help="Path to files to compress",
|
306 |
+
)
|
307 |
+
parser.add_argument(
|
308 |
+
"--delete",
|
309 |
+
action="store_true",
|
310 |
+
help="A flag that triggers deletion of a single file instead of compression",
|
311 |
+
)
|
312 |
+
parser.add_argument(
|
313 |
+
"-r",
|
314 |
+
action="store_true",
|
315 |
+
help="A flag that triggers recursive search on all subdirectories",
|
316 |
+
)
|
317 |
+
parser.add_argument(
|
318 |
+
"--recursive",
|
319 |
+
action="store_true",
|
320 |
+
help="A flag that triggers recursive search on all subdirectories",
|
321 |
+
)
|
322 |
+
parser.add_argument(
|
323 |
+
"--force",
|
324 |
+
action="store_true",
|
325 |
+
help="A flag that forces overwriting when compressing.",
|
326 |
+
)
|
327 |
+
parser.add_argument(
|
328 |
+
"--max_processes",
|
329 |
+
type=int,
|
330 |
+
help="The amount of maximum processes.",
|
331 |
+
)
|
332 |
+
parser.add_argument(
|
333 |
+
"--hf_cache",
|
334 |
+
action="store_true",
|
335 |
+
help="A flag that indicates if the file is in the Hugging Face cache. Must either specify --model or --path to the model's snapshot cache.",
|
336 |
+
)
|
337 |
+
parser.add_argument(
|
338 |
+
"--model",
|
339 |
+
type=str,
|
340 |
+
help="Only when using --hf_cache, specify the model name or path. E.g. 'ibm-granite/granite-7b-instruct'",
|
341 |
+
)
|
342 |
+
parser.add_argument(
|
343 |
+
"--model_branch",
|
344 |
+
type=str,
|
345 |
+
default="main",
|
346 |
+
help="Only when using --model, specify the model branch. Default is 'main'",
|
347 |
+
)
|
348 |
+
args = parser.parse_args()
|
349 |
+
optional_kwargs = {}
|
350 |
+
if args.float32:
|
351 |
+
optional_kwargs["dtype"] = 32
|
352 |
+
if args.streaming_chunk_size is not None:
|
353 |
+
optional_kwargs[
|
354 |
+
"streaming_chunk_size"
|
355 |
+
] = args.streaming_chunk_size
|
356 |
+
if args.path is not None:
|
357 |
+
optional_kwargs["path"] = args.path
|
358 |
+
if args.delete:
|
359 |
+
optional_kwargs["delete"] = args.delete
|
360 |
+
if args.r or args.recursive:
|
361 |
+
optional_kwargs["r"] = args.r
|
362 |
+
if args.force:
|
363 |
+
optional_kwargs["force"] = args.force
|
364 |
+
if args.max_processes:
|
365 |
+
optional_kwargs["max_processes"] = (
|
366 |
+
args.max_processes
|
367 |
+
)
|
368 |
+
if args.hf_cache:
|
369 |
+
optional_kwargs["hf_cache"] = args.hf_cache
|
370 |
+
if args.model:
|
371 |
+
optional_kwargs["model"] = args.model
|
372 |
+
if args.model_branch:
|
373 |
+
optional_kwargs[
|
374 |
+
"branch"
|
375 |
+
] = args.model_branch
|
376 |
+
|
377 |
+
check_and_install_zipnn()
|
378 |
+
compress_files_with_suffix(
|
379 |
+
args.suffix, **optional_kwargs
|
380 |
+
)
|
zipnn_decompress_file.py
ADDED
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import subprocess
|
3 |
+
import sys
|
4 |
+
import argparse
|
5 |
+
|
6 |
+
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
7 |
+
|
8 |
+
RED = "\033[91m"
|
9 |
+
YELLOW = "\033[93m"
|
10 |
+
GREEN = "\033[92m"
|
11 |
+
RESET = "\033[0m"
|
12 |
+
|
13 |
+
|
14 |
+
def check_and_install_zipnn():
|
15 |
+
try:
|
16 |
+
import zipnn
|
17 |
+
except ImportError:
|
18 |
+
print("zipnn not found. Installing...")
|
19 |
+
subprocess.check_call([sys.executable, "-m", "pip", "install", "zipnn"])
|
20 |
+
import zipnn
|
21 |
+
|
22 |
+
|
23 |
+
def decompress_file(input_file, delete=False, force=False, hf_cache=False):
|
24 |
+
import zipnn
|
25 |
+
|
26 |
+
if not input_file.endswith(".znn"):
|
27 |
+
raise ValueError("Input file does not have the '.znn' suffix")
|
28 |
+
|
29 |
+
if os.path.exists(input_file):
|
30 |
+
decompressed_path = input_file[:-4]
|
31 |
+
if not force and os.path.exists(decompressed_path):
|
32 |
+
|
33 |
+
user_input = (
|
34 |
+
input(f"{decompressed_path} already exists; overwrite (y/n)? ").strip().lower()
|
35 |
+
)
|
36 |
+
|
37 |
+
if user_input not in ("yes", "y"):
|
38 |
+
print(f"Skipping {input_file}...")
|
39 |
+
return
|
40 |
+
print(f"Decompressing {input_file}...")
|
41 |
+
|
42 |
+
output_file = input_file[:-4]
|
43 |
+
zpn = zipnn.ZipNN(is_streaming=True)
|
44 |
+
|
45 |
+
with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
|
46 |
+
d_data = b""
|
47 |
+
chunk = infile.read()
|
48 |
+
d_data += zpn.decompress(chunk)
|
49 |
+
outfile.write(d_data)
|
50 |
+
print(f"Decompressed {input_file} to {output_file}")
|
51 |
+
|
52 |
+
if delete and not hf_cache:
|
53 |
+
print(f"Deleting {input_file}...")
|
54 |
+
os.remove(input_file)
|
55 |
+
|
56 |
+
if hf_cache:
|
57 |
+
# If the file is in the Hugging Face cache, fix the symlinks
|
58 |
+
print(f"{YELLOW}Reorganizing Hugging Face cache...{RESET}")
|
59 |
+
try:
|
60 |
+
snapshot_path = os.path.dirname(input_file)
|
61 |
+
blob_name = os.path.join(snapshot_path, os.readlink(input_file))
|
62 |
+
os.rename(output_file, blob_name)
|
63 |
+
os.symlink(blob_name, output_file)
|
64 |
+
|
65 |
+
if os.path.exists(input_file):
|
66 |
+
os.remove(input_file)
|
67 |
+
except Exception as e:
|
68 |
+
raise Exception(f"Error reorganizing Hugging Face cache: {e}")
|
69 |
+
|
70 |
+
else:
|
71 |
+
print(f"Error: The file {input_file} does not exist.")
|
72 |
+
|
73 |
+
|
74 |
+
if __name__ == "__main__":
|
75 |
+
check_and_install_zipnn()
|
76 |
+
|
77 |
+
parser = argparse.ArgumentParser(description="Enter a file path to decompress.")
|
78 |
+
parser.add_argument("input_file", type=str, help="Specify the path to the file to decompress.")
|
79 |
+
parser.add_argument(
|
80 |
+
"--delete",
|
81 |
+
action="store_true",
|
82 |
+
help="A flag that triggers deletion of a single compressed file instead of decompression",
|
83 |
+
)
|
84 |
+
parser.add_argument(
|
85 |
+
"--force", action="store_true", help="A flag that forces overwriting when decompressing."
|
86 |
+
)
|
87 |
+
parser.add_argument(
|
88 |
+
"--hf_cache",
|
89 |
+
action="store_true",
|
90 |
+
help="A flag that indicates if the file is in the Hugging Face cache.",
|
91 |
+
)
|
92 |
+
args = parser.parse_args()
|
93 |
+
optional_kwargs = {}
|
94 |
+
if args.delete:
|
95 |
+
optional_kwargs["delete"] = args.delete
|
96 |
+
if args.force:
|
97 |
+
optional_kwargs["force"] = args.force
|
98 |
+
if args.hf_cache:
|
99 |
+
optional_kwargs["hf_cache"] = args.hf_cache
|
100 |
+
|
101 |
+
decompress_file(args.input_file, **optional_kwargs)
|
zipnn_decompress_path.py
ADDED
@@ -0,0 +1,303 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import sys
|
3 |
+
import argparse
|
4 |
+
import subprocess
|
5 |
+
from pathlib import Path
|
6 |
+
from concurrent.futures import (
|
7 |
+
ProcessPoolExecutor,
|
8 |
+
as_completed,
|
9 |
+
)
|
10 |
+
from zipnn_decompress_file import (
|
11 |
+
decompress_file,
|
12 |
+
)
|
13 |
+
|
14 |
+
sys.path.append(
|
15 |
+
os.path.abspath(
|
16 |
+
os.path.join(
|
17 |
+
os.path.dirname(__file__),
|
18 |
+
"..",
|
19 |
+
)
|
20 |
+
)
|
21 |
+
)
|
22 |
+
|
23 |
+
RED = "\033[91m"
|
24 |
+
YELLOW = "\033[93m"
|
25 |
+
GREEN = "\033[92m"
|
26 |
+
RESET = "\033[0m"
|
27 |
+
|
28 |
+
def check_and_install_zipnn():
|
29 |
+
try:
|
30 |
+
import zipnn
|
31 |
+
except ImportError:
|
32 |
+
print("zipnn not found. Installing...")
|
33 |
+
subprocess.check_call(
|
34 |
+
[
|
35 |
+
sys.executable,
|
36 |
+
"-m",
|
37 |
+
"pip",
|
38 |
+
"install",
|
39 |
+
"zipnn",
|
40 |
+
]
|
41 |
+
)
|
42 |
+
import zipnn
|
43 |
+
|
44 |
+
def replace_in_file(file_path, old: str, new: str) -> None:
|
45 |
+
"""Given a file_path, replace all occurrences of `old` with `new` inpalce."""
|
46 |
+
|
47 |
+
with open(file_path, 'r') as file:
|
48 |
+
file_data = file.read()
|
49 |
+
|
50 |
+
file_data = file_data.replace(old, new)
|
51 |
+
|
52 |
+
with open(file_path, 'w') as file:
|
53 |
+
file.write(file_data)
|
54 |
+
|
55 |
+
def decompress_znn_files(
|
56 |
+
path=".",
|
57 |
+
delete=False,
|
58 |
+
force=False,
|
59 |
+
max_processes=1,
|
60 |
+
hf_cache=False,
|
61 |
+
model="",
|
62 |
+
branch="main",
|
63 |
+
):
|
64 |
+
import zipnn
|
65 |
+
|
66 |
+
overwrite_first=True
|
67 |
+
|
68 |
+
if model:
|
69 |
+
if not hf_cache:
|
70 |
+
raise ValueError(
|
71 |
+
"Must specify --hf_cache when using --model"
|
72 |
+
)
|
73 |
+
try:
|
74 |
+
from huggingface_hub import scan_cache_dir
|
75 |
+
except ImportError:
|
76 |
+
raise ImportError(
|
77 |
+
"huggingface_hub not found. Please pip install huggingface_hub."
|
78 |
+
)
|
79 |
+
cache = scan_cache_dir()
|
80 |
+
repo = next((repo for repo in cache.repos if repo.repo_id == model), None)
|
81 |
+
|
82 |
+
if repo is not None:
|
83 |
+
print(f"Found repo {model} in cache")
|
84 |
+
|
85 |
+
# Get the latest revision path
|
86 |
+
hash = ''
|
87 |
+
try:
|
88 |
+
with open(os.path.join(repo.repo_path, 'refs', branch), "r") as ref:
|
89 |
+
hash = ref.read()
|
90 |
+
except FileNotFoundError:
|
91 |
+
raise FileNotFoundError(f"Branch {branch} not found in repo {model}")
|
92 |
+
|
93 |
+
path = os.path.join(repo.repo_path, 'snapshots', hash)
|
94 |
+
|
95 |
+
file_list = []
|
96 |
+
directories_to_search = [
|
97 |
+
(
|
98 |
+
path,
|
99 |
+
[],
|
100 |
+
os.listdir(path),
|
101 |
+
)
|
102 |
+
]
|
103 |
+
for (
|
104 |
+
root,
|
105 |
+
_,
|
106 |
+
files,
|
107 |
+
) in directories_to_search:
|
108 |
+
for file_name in files:
|
109 |
+
if file_name.endswith(".znn"):
|
110 |
+
decompressed_path = file_name[:-4]
|
111 |
+
if not force and os.path.exists(
|
112 |
+
decompressed_path
|
113 |
+
):
|
114 |
+
#
|
115 |
+
if overwrite_first:
|
116 |
+
overwrite_first=False
|
117 |
+
user_input = (
|
118 |
+
input(
|
119 |
+
f"Decompressed files already exists; Would you like to overwrite them all (y/n)? "
|
120 |
+
)
|
121 |
+
.strip()
|
122 |
+
.lower()
|
123 |
+
)
|
124 |
+
if user_input not in (
|
125 |
+
"y",
|
126 |
+
"yes",
|
127 |
+
):
|
128 |
+
print(
|
129 |
+
f"No forced overwriting."
|
130 |
+
)
|
131 |
+
else:
|
132 |
+
print(
|
133 |
+
f"Overwriting all decompressed files."
|
134 |
+
)
|
135 |
+
force=True
|
136 |
+
|
137 |
+
#
|
138 |
+
if not force:
|
139 |
+
user_input = (
|
140 |
+
input(
|
141 |
+
f"{decompressed_path} already exists; overwrite (y/n)? "
|
142 |
+
)
|
143 |
+
.strip()
|
144 |
+
.lower()
|
145 |
+
)
|
146 |
+
if user_input not in (
|
147 |
+
"y",
|
148 |
+
"yes",
|
149 |
+
):
|
150 |
+
print(
|
151 |
+
f"Skipping {file_name}..."
|
152 |
+
)
|
153 |
+
continue
|
154 |
+
full_path = os.path.join(
|
155 |
+
root,
|
156 |
+
file_name,
|
157 |
+
)
|
158 |
+
file_list.append(full_path)
|
159 |
+
|
160 |
+
if file_list and hf_cache:
|
161 |
+
try:
|
162 |
+
from transformers.utils import (
|
163 |
+
SAFE_WEIGHTS_INDEX_NAME,
|
164 |
+
WEIGHTS_INDEX_NAME
|
165 |
+
)
|
166 |
+
except ImportError:
|
167 |
+
raise ImportError(
|
168 |
+
"Transformers not found. Please pip install transformers."
|
169 |
+
)
|
170 |
+
|
171 |
+
suffix = file_list[0].split('/')[-1].split('.')[-2] # get the one before .znn
|
172 |
+
|
173 |
+
if os.path.exists(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)):
|
174 |
+
print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
|
175 |
+
blob_name = os.path.join(path, os.readlink(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)))
|
176 |
+
replace_in_file(
|
177 |
+
file_path=blob_name,
|
178 |
+
old=f"{suffix}.znn",
|
179 |
+
new=f"{suffix}"
|
180 |
+
)
|
181 |
+
elif os.path.exists(os.path.join(path, WEIGHTS_INDEX_NAME)):
|
182 |
+
print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
|
183 |
+
blob_name = os.path.join(path, os.readlink(os.path.join(path, WEIGHTS_INDEX_NAME)))
|
184 |
+
replace_in_file(
|
185 |
+
file_path=blob_name,
|
186 |
+
old=f"{suffix}.znn",
|
187 |
+
new=f"{suffix}"
|
188 |
+
)
|
189 |
+
|
190 |
+
with ProcessPoolExecutor(
|
191 |
+
max_workers=max_processes
|
192 |
+
) as executor:
|
193 |
+
for file in file_list[:max_processes]:
|
194 |
+
future_to_file = {
|
195 |
+
executor.submit(
|
196 |
+
decompress_file,
|
197 |
+
file,
|
198 |
+
delete,
|
199 |
+
True,
|
200 |
+
hf_cache,
|
201 |
+
): file
|
202 |
+
for file in file_list[
|
203 |
+
:max_processes
|
204 |
+
]
|
205 |
+
}
|
206 |
+
|
207 |
+
file_list = file_list[max_processes:]
|
208 |
+
while future_to_file:
|
209 |
+
|
210 |
+
for future in as_completed(
|
211 |
+
future_to_file
|
212 |
+
):
|
213 |
+
file = future_to_file.pop(
|
214 |
+
future
|
215 |
+
)
|
216 |
+
try:
|
217 |
+
future.result()
|
218 |
+
except Exception as exc:
|
219 |
+
print(
|
220 |
+
f"{RED}File {file} generated an exception: {exc}{RESET}"
|
221 |
+
)
|
222 |
+
|
223 |
+
if file_list:
|
224 |
+
next_file = file_list.pop(
|
225 |
+
0
|
226 |
+
)
|
227 |
+
future_to_file[
|
228 |
+
executor.submit(
|
229 |
+
decompress_file,
|
230 |
+
next_file,
|
231 |
+
delete,
|
232 |
+
True,
|
233 |
+
hf_cache,
|
234 |
+
)
|
235 |
+
] = next_file
|
236 |
+
#
|
237 |
+
print(f"{GREEN}All files decompressed{RESET}")
|
238 |
+
|
239 |
+
|
240 |
+
if __name__ == "__main__":
|
241 |
+
check_and_install_zipnn()
|
242 |
+
|
243 |
+
parser = argparse.ArgumentParser(
|
244 |
+
description="Compresses all .znn files."
|
245 |
+
)
|
246 |
+
parser.add_argument(
|
247 |
+
"--path",
|
248 |
+
type=str,
|
249 |
+
help="Path to folder of files to decompress. If left empty, checks current folder.",
|
250 |
+
)
|
251 |
+
parser.add_argument(
|
252 |
+
"--delete",
|
253 |
+
action="store_true",
|
254 |
+
help="A flag that triggers deletion of a single compressed file instead of decompression",
|
255 |
+
)
|
256 |
+
parser.add_argument(
|
257 |
+
"--force",
|
258 |
+
action="store_true",
|
259 |
+
help="A flag that forces overwriting when decompressing.",
|
260 |
+
)
|
261 |
+
parser.add_argument(
|
262 |
+
"--max_processes",
|
263 |
+
type=int,
|
264 |
+
help="The amount of maximum processes.",
|
265 |
+
)
|
266 |
+
parser.add_argument(
|
267 |
+
"--hf_cache",
|
268 |
+
action="store_true",
|
269 |
+
help="A flag that indicates if the file is in the Hugging Face cache. Must either specify --model or --path to the model's snapshot cache.",
|
270 |
+
)
|
271 |
+
parser.add_argument(
|
272 |
+
"--model",
|
273 |
+
type=str,
|
274 |
+
help="Only when using --hf_cache, specify the model name or path. E.g. 'ibm-granite/granite-7b-instruct'",
|
275 |
+
)
|
276 |
+
parser.add_argument(
|
277 |
+
"--model_branch",
|
278 |
+
type=str,
|
279 |
+
default="main",
|
280 |
+
help="Only when using --model, specify the model branch. Default is 'main'",
|
281 |
+
)
|
282 |
+
args = parser.parse_args()
|
283 |
+
optional_kwargs = {}
|
284 |
+
if args.path is not None:
|
285 |
+
optional_kwargs["path"] = args.path
|
286 |
+
if args.delete:
|
287 |
+
optional_kwargs["delete"] = args.delete
|
288 |
+
if args.force:
|
289 |
+
optional_kwargs["force"] = args.force
|
290 |
+
if args.max_processes:
|
291 |
+
optional_kwargs["max_processes"] = (
|
292 |
+
args.max_processes
|
293 |
+
)
|
294 |
+
if args.hf_cache:
|
295 |
+
optional_kwargs["hf_cache"] = args.hf_cache
|
296 |
+
if args.model:
|
297 |
+
optional_kwargs["model"] = args.model
|
298 |
+
if args.model_branch:
|
299 |
+
optional_kwargs[
|
300 |
+
"branch"
|
301 |
+
] = args.model_branch
|
302 |
+
|
303 |
+
decompress_znn_files(**optional_kwargs)
|