royleibov commited on
Commit
a99e639
•
1 Parent(s): 3ca981c

Add .znn files

Browse files
.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:eab4f4dc1abf860794c98ce3759b4ace1059fc0fc041ede76b88988b3557c132
3
- size 3898649896
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f7552c637cf4bb0a59c40215e084bc5a8cbbdabefce3da42db774c84ddb4b02c
3
+ size 2620478673
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:0546b7cd070ff239de0f61a28770de6a0ac517ad313071e151653e38911ffb25
3
- size 3864726320
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2a0586c7d125105455df858b3aedd4ebcf4fea0a1854b0b6b3dc74c4da0ce0c5
3
+ size 2565425513
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:11368ea1e9c0039ca54012a74ac32d543a9e9d980b2d0ad7e93af0341ace1fe6
3
- size 3864726424
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:86b0e141b4b405babec9abcfadcc6b645ca320773f78fc067e4b4b1daaeeab77
3
+ size 2565999485
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:381e110d74db594cc24956af0e74a01ae02104497113a78926990c5bd1f395f6
3
- size 3864733680
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:762df4e13c1a6465f0b2b69261dbd806cb81bcb75364abfb7177e6f676c2e37e
3
+ size 2564782431
model-00005-of-00005.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:e76ec4dc2e7c0dcbe3432e58fae08e62204a2ef2323069d85c8fa0affb87332a
3
- size 1089994880
 
 
 
 
model-00005-of-00005.safetensors.znn ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3782e236ec2c3d412e4ebf6d4008d87e20809e2ae429917992191fbec3cd1a9b
3
+ size 725050380
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,178 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ if delete and not hf_cache:
73
+ print(f"Deleting {full_path}...")
74
+ os.remove(full_path)
75
+ else:
76
+ compressed_path = full_path + ".znn"
77
+ if not force and os.path.exists(compressed_path):
78
+ user_input = (
79
+ input(f"{compressed_path} already exists; overwrite (y/n)? ").strip().lower()
80
+ )
81
+ if user_input not in ("yes", "y"):
82
+ print(f"Skipping {full_path}...")
83
+ return
84
+ print(f"Compressing {full_path}...")
85
+ #
86
+ output_file = input_file + ".znn"
87
+ if dtype:
88
+ zpn = zipnn.ZipNN(
89
+ bytearray_dtype="float32",
90
+ is_streaming=True,
91
+ streaming_chunk_kb=streaming_chunk_size,
92
+ )
93
+ else:
94
+ zpn = zipnn.ZipNN(
95
+ is_streaming=True,
96
+ streaming_chunk_kb=streaming_chunk_size,
97
+ )
98
+ file_size_before = 0
99
+ file_size_after = 0
100
+ start_time = time.time()
101
+ with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
102
+ chunk = infile.read()
103
+ file_size_before += len(chunk)
104
+ compressed_chunk = zpn.compress(chunk)
105
+ if compressed_chunk:
106
+ file_size_after += len(compressed_chunk)
107
+ outfile.write(compressed_chunk)
108
+ end_time = time.time() - start_time
109
+ print(f"Compressed {input_file} to {output_file}")
110
+ print(
111
+ 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}"
112
+ )
113
+
114
+ if hf_cache:
115
+ # If the file is in the Hugging Face cache, fix the symlinks
116
+ print(f"{YELLOW}Reorganizing Hugging Face cache...{RESET}")
117
+ try:
118
+ snapshot_path = os.path.dirname(input_file)
119
+ blob_name = os.path.join(snapshot_path, os.readlink(input_file))
120
+ os.rename(output_file, blob_name)
121
+ os.symlink(blob_name, output_file)
122
+ if os.path.exists(input_file):
123
+ os.remove(input_file)
124
+ except Exception as e:
125
+ raise Exception(f"Error reorganizing Hugging Face cache: {e}")
126
+
127
+ if __name__ == "__main__":
128
+ if len(sys.argv) < 2:
129
+ print("Usage: python compress_files.py <suffix>")
130
+ print("Example: python compress_files.py 'safetensors'")
131
+ sys.exit(1)
132
+
133
+ parser = argparse.ArgumentParser(description="Enter a file path to compress.")
134
+ parser.add_argument(
135
+ "input_file",
136
+ type=str,
137
+ help="Specify the path to the file to compress.",
138
+ )
139
+ parser.add_argument(
140
+ "--float32",
141
+ action="store_true",
142
+ help="A flag that triggers float32 compression",
143
+ )
144
+ parser.add_argument(
145
+ "--streaming_chunk_size",
146
+ type=str,
147
+ help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
148
+ )
149
+ parser.add_argument(
150
+ "--delete",
151
+ action="store_true",
152
+ help="A flag that triggers deletion of a single file instead of compression",
153
+ )
154
+ parser.add_argument(
155
+ "--force",
156
+ action="store_true",
157
+ help="A flag that forces overwriting when compressing.",
158
+ )
159
+ parser.add_argument(
160
+ "--hf_cache",
161
+ action="store_true",
162
+ help="A flag that indicates if the file is in the Hugging Face cache.",
163
+ )
164
+ args = parser.parse_args()
165
+ optional_kwargs = {}
166
+ if args.float32:
167
+ optional_kwargs["dtype"] = 32
168
+ if args.streaming_chunk_size is not None:
169
+ optional_kwargs["streaming_chunk_size"] = args.streaming_chunk_size
170
+ if args.delete:
171
+ optional_kwargs["delete"] = args.delete
172
+ if args.force:
173
+ optional_kwargs["force"] = args.force
174
+ if args.hf_cache:
175
+ optional_kwargs["hf_cache"] = args.hf_cache
176
+
177
+ check_and_install_zipnn()
178
+ 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: Path | str, 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
+ if delete and not hf_cache:
31
+ print(f"Deleting {input_file}...")
32
+ os.remove(input_file)
33
+ else:
34
+ decompressed_path = input_file[:-4]
35
+ if not force and os.path.exists(decompressed_path):
36
+
37
+ user_input = (
38
+ input(f"{decompressed_path} already exists; overwrite (y/n)? ").strip().lower()
39
+ )
40
+
41
+ if user_input not in ("yes", "y"):
42
+ print(f"Skipping {input_file}...")
43
+ return
44
+ print(f"Decompressing {input_file}...")
45
+
46
+ output_file = input_file[:-4]
47
+ zpn = zipnn.ZipNN(is_streaming=True)
48
+
49
+ with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
50
+ d_data = b""
51
+ chunk = infile.read()
52
+ d_data += zpn.decompress(chunk)
53
+ outfile.write(d_data)
54
+ print(f"Decompressed {input_file} to {output_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: Path | str, 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)