add znn

.gitattributes

@@ -8,3 +8,4 @@
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
 model.safetensors filter=lfs diff=lfs merge=lfs -text
+*.znn filter=lfs diff=lfs merge=lfs -text

model.safetensors → model.safetensors.znn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5bde1d28afb363d0103324efeb5afc8b2b397fe5e04beabb9b1ef355255ade81
-size 498818054
+oid sha256:173019f480339b3e3bf8cae56a0e67b099d1b8500a178121c34db2456d630178
+size 273378294

zipnn_compress_file.py

@@ -0,0 +1,179 @@
+import os
+import subprocess
+import sys
+import argparse
+import time
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+KB = 1024
+MB = 1024 * 1024
+GB = 1024 * 1024 * 1024
+
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "zipnn",
+                "--upgrade",
+            ]
+        )
+        import zipnn
+
+
+def parse_streaming_chunk_size(
+    streaming_chunk_size,
+):
+    if str(streaming_chunk_size).isdigit():
+        final = int(streaming_chunk_size)
+    else:
+        size_value = int(streaming_chunk_size[:-2])
+        size_unit = streaming_chunk_size[-2].lower()
+
+        if size_unit == "k":
+            final = KB * size_value
+        elif size_unit == "m":
+            final = MB * size_value
+        elif size_unit == "g":
+            final = GB * size_value
+        else:
+            raise ValueError(f"Invalid size unit: {size_unit}. Use 'k', 'm', or 'g'.")
+
+    return final
+
+
+def compress_file(
+    input_file,
+    dtype="",
+    streaming_chunk_size=1048576,
+    delete=False,
+    force=False,
+    hf_cache=False,
+):
+    import zipnn
+
+    streaming_chunk_size = parse_streaming_chunk_size(streaming_chunk_size)
+    full_path = input_file
+    if not os.path.exists(full_path):
+        print(f"{RED}File not found{RESET}")
+        return
+
+    compressed_path = full_path + ".znn"
+    if not force and os.path.exists(compressed_path):
+        user_input = (
+            input(f"{compressed_path} already exists; overwrite (y/n)? ").strip().lower()
+        )
+        if user_input not in ("yes", "y"):
+            print(f"Skipping {full_path}...")
+            return
+    print(f"Compressing {full_path}...")
+    #
+    output_file = input_file + ".znn"
+    if dtype:
+        zpn = zipnn.ZipNN(
+            bytearray_dtype="float32",
+            is_streaming=True,
+            streaming_chunk_kb=streaming_chunk_size,
+        )
+    else:
+        zpn = zipnn.ZipNN(
+            is_streaming=True,
+            streaming_chunk_kb=streaming_chunk_size,
+        )
+    file_size_before = 0
+    file_size_after = 0
+    start_time = time.time()
+    with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
+        chunk = infile.read()
+        file_size_before += len(chunk)
+        compressed_chunk = zpn.compress(chunk)
+        if compressed_chunk:
+            file_size_after += len(compressed_chunk)
+            outfile.write(compressed_chunk)
+    end_time = time.time() - start_time
+    print(f"Compressed {input_file} to {output_file}")
+    print(
+        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}"
+    )
+
+    if delete and not hf_cache:
+        print(f"Deleting {full_path}...")
+        os.remove(full_path)
+
+    if hf_cache:
+        # If the file is in the Hugging Face cache, fix the symlinks
+        print(f"{YELLOW}Reorganizing Hugging Face cache...{RESET}")
+        try:
+            snapshot_path = os.path.dirname(input_file)
+            blob_name = os.path.join(snapshot_path, os.readlink(input_file))
+            os.rename(output_file, blob_name)
+            os.symlink(blob_name, output_file)
+            if os.path.exists(input_file):
+                os.remove(input_file)
+        except Exception as e:
+            raise Exception(f"Error reorganizing Hugging Face cache: {e}")
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python compress_files.py <suffix>")
+        print("Example: python compress_files.py 'safetensors'")
+        sys.exit(1)
+
+    parser = argparse.ArgumentParser(description="Enter a file path to compress.")
+    parser.add_argument(
+        "input_file",
+        type=str,
+        help="Specify the path to the file to compress.",
+    )
+    parser.add_argument(
+        "--float32",
+        action="store_true",
+        help="A flag that triggers float32 compression",
+    )
+    parser.add_argument(
+        "--streaming_chunk_size",
+        type=str,
+        help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single file instead of compression",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="A flag that forces overwriting when compressing.",
+    )
+    parser.add_argument(
+        "--hf_cache",
+        action="store_true",
+        help="A flag that indicates if the file is in the Hugging Face cache.",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.streaming_chunk_size is not None:
+        optional_kwargs["streaming_chunk_size"] = args.streaming_chunk_size
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+    if args.hf_cache:
+        optional_kwargs["hf_cache"] = args.hf_cache
+
+    check_and_install_zipnn()
+    compress_file(args.input_file, **optional_kwargs)

zipnn_compress_path.py

@@ -0,0 +1,380 @@
+import os
+import subprocess
+import sys
+import argparse
+from pathlib import Path
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    as_completed,
+)
+from zipnn_compress_file import compress_file
+
+sys.path.append(
+    os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__), ".."
+        )
+    )
+)
+
+
+KB = 1024
+MB = 1024 * 1024
+GB = 1024 * 1024 * 1024
+
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "zipnn",
+                "--upgrade",
+            ]
+        )
+        import zipnn
+
+
+def parse_streaming_chunk_size(
+    streaming_chunk_size,
+):
+    if str(streaming_chunk_size).isdigit():
+        final = int(streaming_chunk_size)
+    else:
+        size_value = int(
+            streaming_chunk_size[:-2]
+        )
+        size_unit = streaming_chunk_size[
+            -2
+        ].lower()
+
+        if size_unit == "k":
+            final = KB * size_value
+        elif size_unit == "m":
+            final = MB * size_value
+        elif size_unit == "g":
+            final = GB * size_value
+        else:
+            raise ValueError(
+                f"Invalid size unit: {size_unit}. Use 'k', 'm', or 'g'."
+            )
+
+    return final
+
+def replace_in_file(file_path, old: str, new: str) -> None:
+    """Given a file_path, replace all occurrences of `old` with `new` inpalce."""
+
+    with open(file_path, 'r') as file:
+        file_data = file.read()
+
+    file_data = file_data.replace(old, new)
+
+    with open(file_path, 'w') as file:
+        file.write(file_data)
+
+def compress_files_with_suffix(
+    suffix,
+    dtype="",
+    streaming_chunk_size=1048576,
+    path=".",
+    delete=False,
+    r=False,
+    force=False,
+    max_processes=1,
+    hf_cache=False,
+    model="",
+    branch="main",
+):
+    import zipnn
+
+    overwrite_first=True
+    file_list = []
+    streaming_chunk_size = (
+        parse_streaming_chunk_size(
+            streaming_chunk_size
+        )
+    )
+    if model:
+        if not hf_cache:
+            raise ValueError(
+                "Must specify --hf_cache when using --model"
+            )
+        try:
+            from huggingface_hub import scan_cache_dir
+        except ImportError:
+            raise ImportError(
+                "huggingface_hub not found. Please pip install huggingface_hub."
+            )  
+        cache = scan_cache_dir()
+        repo = next((repo for repo in cache.repos if repo.repo_id == model), None)
+
+        if repo is not None:
+            print(f"Found repo {model} in cache")
+            
+            # Get the latest revision path
+            hash = ''
+            try:
+                with open(os.path.join(repo.repo_path, 'refs', branch), "r") as ref:
+                    hash = ref.read()
+            except FileNotFoundError:
+                raise FileNotFoundError(f"Branch {branch} not found in repo {model}")
+            
+            path = os.path.join(repo.repo_path, 'snapshots', hash)
+
+    directories_to_search = (
+        os.walk(path)
+        if r
+        else [(path, [], os.listdir(path))]
+    )
+    files_found = False
+    for root, _, files in directories_to_search:
+        for file_name in files:
+            if file_name.endswith(suffix):
+                compressed_path = (
+                    file_name + ".znn"
+                )
+                if not force and os.path.exists(
+                    compressed_path
+                ):
+                    #
+                    if overwrite_first:
+                        overwrite_first=False
+                        user_input = (
+                            input(
+                                f"Compressed files already exists; Would you like to overwrite them all (y/n)? "
+                            )
+                            .strip()
+                            .lower()
+                        )
+                        if user_input not in (
+                            "y",
+                            "yes",
+                        ):
+                            print(
+                                f"No forced overwriting."
+                            )
+                        else:
+                            print(
+                                f"Overwriting all compressed files."
+                            )
+                            force=True
+                    #
+                    if not force:
+                        user_input = (
+                            input(
+                                f"{compressed_path} already exists; overwrite (y/n)? "
+                            )
+                            .strip()
+                            .lower()
+                        )
+                        if user_input not in (
+                            "y",
+                            "yes",
+                        ):
+                            print(
+                                f"Skipping {file_name}..."
+                            )
+                            continue
+                files_found = True
+                full_path = os.path.join(
+                    root, file_name
+                )
+                file_list.append(full_path)
+
+    if file_list and hf_cache:
+        try:
+            from transformers.utils import (
+                SAFE_WEIGHTS_INDEX_NAME,
+                WEIGHTS_INDEX_NAME
+            )
+        except ImportError:
+            raise ImportError(
+                "Transformers not found. Please pip install transformers."
+            )
+        
+        if os.path.exists(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)):
+            print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
+            blob_name = os.path.join(path, os.readlink(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)))
+            replace_in_file(
+                    file_path=blob_name,
+                    old=f"{suffix}",
+                    new=f"{suffix}.znn"
+                )
+        elif os.path.exists(os.path.join(path, WEIGHTS_INDEX_NAME)):
+            print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
+            blob_name = os.path.join(path, os.readlink(os.path.join(path, WEIGHTS_INDEX_NAME)))
+            replace_in_file(
+                    file_path=blob_name,
+                    old=f"{suffix}",
+                    new=f"{suffix}.znn"
+                )
+
+    with ProcessPoolExecutor(
+        max_workers=max_processes
+    ) as executor:
+        future_to_file = {
+            executor.submit(
+                compress_file,
+                file,
+                dtype,
+                streaming_chunk_size,
+                delete,
+                True,
+                hf_cache,
+            ): file
+            for file in file_list[:max_processes]
+        }
+        file_list = file_list[max_processes:]
+        while future_to_file:
+            for future in as_completed(
+                future_to_file
+            ):
+                file = future_to_file.pop(future)
+
+                try:
+                    future.result()
+                except Exception as exc:
+                    print(
+                        f"{RED}File {file} generated an exception: {exc}{RESET}"
+                    )
+
+                if file_list:
+                    next_file = file_list.pop(0)
+                    future_to_file[
+                        executor.submit(
+                            compress_file,
+                            next_file,
+                            dtype,
+                            streaming_chunk_size,
+                            delete,
+                            True,
+                            hf_cache,
+                        )
+                    ] = next_file
+
+    if not files_found:
+        print(
+            f"{RED}No files with the suffix '{suffix}' found.{RESET}"
+        )
+
+    print(f"{GREEN}All files compressed{RESET}")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(
+            "Usage: python compress_files.py <suffix>"
+        )
+        print(
+            "Example: python compress_files.py 'safetensors'"
+        )
+        sys.exit(1)
+
+    parser = argparse.ArgumentParser(
+        description="Enter a suffix to compress, (optional) dtype, (optional) streaming chunk size, (optional) path to files."
+    )
+    parser.add_argument(
+        "suffix",
+        type=str,
+        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.",
+    )
+    parser.add_argument(
+        "--float32",
+        action="store_true",
+        help="A flag that triggers float32 compression",
+    )
+    parser.add_argument(
+        "--streaming_chunk_size",
+        type=str,
+        help="An optional streaming chunk size. The format is int (for size in Bytes) or int+KB/MB/GB. Default is 1MB",
+    )
+    parser.add_argument(
+        "--path",
+        type=str,
+        help="Path to files to compress",
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single file instead of compression",
+    )
+    parser.add_argument(
+        "-r",
+        action="store_true",
+        help="A flag that triggers recursive search on all subdirectories",
+    )
+    parser.add_argument(
+        "--recursive",
+        action="store_true",
+        help="A flag that triggers recursive search on all subdirectories",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="A flag that forces overwriting when compressing.",
+    )
+    parser.add_argument(
+        "--max_processes",
+        type=int,
+        help="The amount of maximum processes.",
+    )
+    parser.add_argument(
+        "--hf_cache",
+        action="store_true",
+        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.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        help="Only when using --hf_cache, specify the model name or path. E.g. 'ibm-granite/granite-7b-instruct'",
+    )
+    parser.add_argument(
+        "--model_branch",
+        type=str,
+        default="main",
+        help="Only when using --model, specify the model branch. Default is 'main'",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.streaming_chunk_size is not None:
+        optional_kwargs[
+            "streaming_chunk_size"
+        ] = args.streaming_chunk_size
+    if args.path is not None:
+        optional_kwargs["path"] = args.path
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.r or args.recursive:
+        optional_kwargs["r"] = args.r
+    if args.force:
+        optional_kwargs["force"] = args.force
+    if args.max_processes:
+        optional_kwargs["max_processes"] = (
+            args.max_processes
+        )
+    if args.hf_cache:
+        optional_kwargs["hf_cache"] = args.hf_cache
+    if args.model:
+        optional_kwargs["model"] = args.model
+    if args.model_branch:
+        optional_kwargs[
+            "branch"
+        ] = args.model_branch
+
+    check_and_install_zipnn()
+    compress_files_with_suffix(
+        args.suffix, **optional_kwargs
+    )

zipnn_decompress_file.py

@@ -0,0 +1,114 @@
+import os
+import subprocess
+import sys
+import argparse
+import time
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+GB = 1024 * 1024 * 1024
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "zipnn"])
+        import zipnn
+
+
+def decompress_file(input_file, delete=False, force=False, hf_cache=False):
+    import zipnn
+
+    if not input_file.endswith(".znn"):
+        raise ValueError("Input file does not have the '.znn' suffix")
+
+    if os.path.exists(input_file):
+        decompressed_path = input_file[:-4]
+        if not force and os.path.exists(decompressed_path):
+
+            user_input = (
+                input(f"{decompressed_path} already exists; overwrite (y/n)? ").strip().lower()
+            )
+
+            if user_input not in ("yes", "y"):
+                print(f"Skipping {input_file}...")
+                return
+        print(f"Decompressing {input_file}...")
+
+        output_file = input_file[:-4]
+        zpn = zipnn.ZipNN(is_streaming=True)
+
+        file_size_before = 0
+        file_size_after = 0
+        start_time = time.time()
+        with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
+            d_data = b""
+            chunk = infile.read()
+            file_size_before = len(chunk)
+            d_data += zpn.decompress(chunk)
+            file_size_after = len(d_data)
+            outfile.write(d_data)
+            print(f"Decompressed {input_file} to {output_file}")
+        end_time = time.time() - start_time
+
+        print(
+            f"{GREEN}Back to original size: {file_size_after/GB:.02f}GB size before decompression: {file_size_before/GB:.02f}GB, time: {end_time:.02f}{RESET}"
+            )
+
+
+        if delete and not hf_cache:
+            print(f"Deleting {input_file}...")
+            os.remove(input_file)
+
+        if hf_cache:
+            # If the file is in the Hugging Face cache, fix the symlinks
+            print(f"{YELLOW}Reorganizing Hugging Face cache...{RESET}")
+            try:
+                snapshot_path = os.path.dirname(input_file)
+                blob_name = os.path.join(snapshot_path, os.readlink(input_file))
+                os.rename(output_file, blob_name)
+                os.symlink(blob_name, output_file)
+                
+                if os.path.exists(input_file):
+                    os.remove(input_file)
+            except Exception as e:
+                raise Exception(f"Error reorganizing Hugging Face cache: {e}")
+
+    else:
+        print(f"Error: The file {input_file} does not exist.")
+
+
+if __name__ == "__main__":
+    check_and_install_zipnn()
+
+    parser = argparse.ArgumentParser(description="Enter a file path to decompress.")
+    parser.add_argument("input_file", type=str, help="Specify the path to the file to decompress.")
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single compressed file instead of decompression",
+    )
+    parser.add_argument(
+        "--force", action="store_true", help="A flag that forces overwriting when decompressing."
+    )
+    parser.add_argument(
+        "--hf_cache",
+        action="store_true",
+        help="A flag that indicates if the file is in the Hugging Face cache.",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+    if args.hf_cache:
+        optional_kwargs["hf_cache"] = args.hf_cache
+
+    decompress_file(args.input_file, **optional_kwargs)

zipnn_decompress_path.py

@@ -0,0 +1,303 @@
+import os
+import sys
+import argparse
+import subprocess
+from pathlib import Path
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    as_completed,
+)
+from zipnn_decompress_file import (
+    decompress_file,
+)
+
+sys.path.append(
+    os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__),
+            "..",
+        )
+    )
+)
+
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+def check_and_install_zipnn():
+    try:
+        import zipnn
+    except ImportError:
+        print("zipnn not found. Installing...")
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "zipnn",
+            ]
+        )
+        import zipnn
+
+def replace_in_file(file_path, old: str, new: str) -> None:
+    """Given a file_path, replace all occurrences of `old` with `new` inpalce."""
+
+    with open(file_path, 'r') as file:
+        file_data = file.read()
+
+    file_data = file_data.replace(old, new)
+
+    with open(file_path, 'w') as file:
+        file.write(file_data)
+
+def decompress_znn_files(
+    path=".",
+    delete=False,
+    force=False,
+    max_processes=1,
+    hf_cache=False,
+    model="",
+    branch="main",
+):
+    import zipnn
+
+    overwrite_first=True
+
+    if model:
+        if not hf_cache:
+            raise ValueError(
+                "Must specify --hf_cache when using --model"
+            )
+        try:
+            from huggingface_hub import scan_cache_dir
+        except ImportError:
+            raise ImportError(
+                "huggingface_hub not found. Please pip install huggingface_hub."
+            )  
+        cache = scan_cache_dir()
+        repo = next((repo for repo in cache.repos if repo.repo_id == model), None)
+
+        if repo is not None:
+            print(f"Found repo {model} in cache")
+            
+            # Get the latest revision path
+            hash = ''
+            try:
+                with open(os.path.join(repo.repo_path, 'refs', branch), "r") as ref:
+                    hash = ref.read()
+            except FileNotFoundError:
+                raise FileNotFoundError(f"Branch {branch} not found in repo {model}")
+            
+            path = os.path.join(repo.repo_path, 'snapshots', hash)
+
+    file_list = []
+    directories_to_search = [
+        (
+            path,
+            [],
+            os.listdir(path),
+        )
+    ]
+    for (
+        root,
+        _,
+        files,
+    ) in directories_to_search:
+        for file_name in files:
+            if file_name.endswith(".znn"):
+                decompressed_path = file_name[:-4]
+                if not force and os.path.exists(
+                    decompressed_path
+                ):
+                    #
+                    if overwrite_first:
+                        overwrite_first=False
+                        user_input = (
+                            input(
+                                f"Decompressed files already exists; Would you like to overwrite them all (y/n)? "
+                            )
+                            .strip()
+                            .lower()
+                        )
+                        if user_input not in (
+                            "y",
+                            "yes",
+                        ):
+                            print(
+                                f"No forced overwriting."
+                            )
+                        else:
+                            print(
+                                f"Overwriting all decompressed files."
+                            )
+                            force=True
+                        
+                    #
+                    if not force:
+                        user_input = (
+                            input(
+                                f"{decompressed_path} already exists; overwrite (y/n)? "
+                            )
+                            .strip()
+                            .lower()
+                        )
+                        if user_input not in (
+                            "y",
+                            "yes",
+                        ):
+                            print(
+                                f"Skipping {file_name}..."
+                            )
+                            continue
+                full_path = os.path.join(
+                    root,
+                    file_name,
+                )
+                file_list.append(full_path)
+
+    if file_list and hf_cache:
+        try:
+            from transformers.utils import (
+                SAFE_WEIGHTS_INDEX_NAME,
+                WEIGHTS_INDEX_NAME
+            )
+        except ImportError:
+            raise ImportError(
+                "Transformers not found. Please pip install transformers."
+            )
+            
+        suffix = file_list[0].split('/')[-1].split('.')[-2] # get the one before .znn
+
+        if os.path.exists(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)):
+            print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
+            blob_name = os.path.join(path, os.readlink(os.path.join(path, SAFE_WEIGHTS_INDEX_NAME)))
+            replace_in_file(
+                    file_path=blob_name,
+                    old=f"{suffix}.znn",
+                    new=f"{suffix}"
+                )
+        elif os.path.exists(os.path.join(path, WEIGHTS_INDEX_NAME)):
+            print(f"{YELLOW}Fixing Hugging Face model json...{RESET}")
+            blob_name = os.path.join(path, os.readlink(os.path.join(path, WEIGHTS_INDEX_NAME)))
+            replace_in_file(
+                    file_path=blob_name,
+                    old=f"{suffix}.znn",
+                    new=f"{suffix}"
+                )
+
+    with ProcessPoolExecutor(
+        max_workers=max_processes
+    ) as executor:
+        for file in file_list[:max_processes]:
+            future_to_file = {
+                executor.submit(
+                    decompress_file,
+                    file,
+                    delete,
+                    True,
+                    hf_cache,
+                ): file
+                for file in file_list[
+                    :max_processes
+                ]
+            }
+
+            file_list = file_list[max_processes:]
+            while future_to_file:
+
+                for future in as_completed(
+                    future_to_file
+                ):
+                    file = future_to_file.pop(
+                        future
+                    )
+                    try:
+                        future.result()
+                    except Exception as exc:
+                        print(
+                            f"{RED}File {file} generated an exception: {exc}{RESET}"
+                        )
+
+                    if file_list:
+                        next_file = file_list.pop(
+                            0
+                        )
+                        future_to_file[
+                            executor.submit(
+                                decompress_file,
+                                next_file,
+                                delete,
+                                True,
+                                hf_cache,
+                            )
+                        ] = next_file
+                        #
+    print(f"{GREEN}All files decompressed{RESET}")
+
+
+if __name__ == "__main__":
+    check_and_install_zipnn()
+
+    parser = argparse.ArgumentParser(
+        description="Compresses all .znn files."
+    )
+    parser.add_argument(
+        "--path",
+        type=str,
+        help="Path to folder of files to decompress. If left empty, checks current folder.",
+    )
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="A flag that triggers deletion of a single compressed file instead of decompression",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="A flag that forces overwriting when decompressing.",
+    )
+    parser.add_argument(
+        "--max_processes",
+        type=int,
+        help="The amount of maximum processes.",
+    )
+    parser.add_argument(
+        "--hf_cache",
+        action="store_true",
+        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.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        help="Only when using --hf_cache, specify the model name or path. E.g. 'ibm-granite/granite-7b-instruct'",
+    )
+    parser.add_argument(
+        "--model_branch",
+        type=str,
+        default="main",
+        help="Only when using --model, specify the model branch. Default is 'main'",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.path is not None:
+        optional_kwargs["path"] = args.path
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+    if args.max_processes:
+        optional_kwargs["max_processes"] = (
+            args.max_processes
+        )
+    if args.hf_cache:
+        optional_kwargs["hf_cache"] = args.hf_cache
+    if args.model:
+        optional_kwargs["model"] = args.model
+    if args.model_branch:
+        optional_kwargs[
+            "branch"
+        ] = args.model_branch
+
+    decompress_znn_files(**optional_kwargs)