Compress with ZipNN

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.znn filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,6 +5,7 @@ tags:
 - jamba
 - mamba
 - moe
+base_model: ai21labs/Jamba-v0.1
 ---
 
 # Model Card for Jamba
@@ -15,6 +16,8 @@ Jamba is the first production-scale Mamba implementation, which opens up interes
 
 This model card is for the base version of Jamba. It’s a pretrained, mixture-of-experts (MoE) generative text model, with 12B active parameters and a total of 52B parameters across all experts. It supports a 256K context length, and can fit up to 140K tokens on a single 80GB GPU.
 
+This fork is compressed using **ZipNN**. To use the model, decompress the model tensors as discribed below and load the **local** weights.
+
 For full details of this model please read the [white paper](https://arxiv.org/abs/2403.19887) and the [release blog post](https://www.ai21.com/blog/announcing-jamba).
 
 ## Model Details
@@ -40,12 +43,25 @@ You also have to have the model on a CUDA device.
 
 You can run the model not using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly lower latencies. In order to do that, you'll need to specify `use_mamba_kernels=False` when loading the model.
 
+You need to [clone this repository](https://huggingface.co/royleibov/Jamba-v0.1-ZipNN-Compressed?clone=true) to decompress the model.
+
+Then:
+```bash
+cd Jamba-v0.1-ZipNN-Compressed
+```
+
+
 ### Run the model
+First decompress the model weights:
+```bash
+python3 zipnn_decompress_path.py --path .
+```
+
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
+model = AutoModelForCausalLM.from_pretrained("PATH_TO_MODEL") # "." if in directory
+tokenizer = AutoTokenizer.from_pretrained("PATH_TO_MODEL") # "." if in directory
 
 input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]
 
@@ -65,7 +81,7 @@ Please note that if you're using `transformers<4.40.0`, `trust_remote_code=True`
 ```python
 from transformers import AutoModelForCausalLM
 import torch
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
+model = AutoModelForCausalLM.from_pretrained("PATH_TO_MODEL", # "." if in directory,
                                              torch_dtype=torch.bfloat16)    # you can also use torch_dtype=torch.float16
 ```
 
@@ -73,7 +89,7 @@ When using half precision, you can enable the [FlashAttention2](https://github.c
 ```python
 from transformers import AutoModelForCausalLM
 import torch
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
+model = AutoModelForCausalLM.from_pretrained("PATH_TO_MODEL", # "." if in directory
                                              torch_dtype=torch.bfloat16,
                                              attn_implementation="flash_attention_2",
                                              device_map="auto")
@@ -88,7 +104,7 @@ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 quantization_config = BitsAndBytesConfig(load_in_8bit=True,
                                          llm_int8_skip_modules=["mamba"])
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
+model = AutoModelForCausalLM.from_pretrained("PATH_TO_MODEL", # "." if in directory
                                              torch_dtype=torch.bfloat16,
                                              attn_implementation="flash_attention_2",
                                              quantization_config=quantization_config)
@@ -105,9 +121,9 @@ from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 
-tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
-model = AutoModelForCausalLM.from_pretrained(
-    "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained("PATH_TO_MODEL") # "." if in directory
+model = AutoModelForCausalLM.from_pretrained("PATH_TO_MODEL", # "." if in directory
+            device_map='auto', torch_dtype=torch.bfloat16)
 
 lora_config = LoraConfig(
     r=8,
@@ -164,4 +180,4 @@ As a base model, Jamba is intended for use as a foundation layer for fine tuning
 ## About AI21
 AI21 builds reliable, practical, and scalable AI solutions for the enterprise.
 
-Jamba is the first in AI21’s new family of models, and the Instruct version of Jamba is coming soon to the [AI21 platform](https://www.ai21.com/studio). 
+Jamba is the first in AI21’s new family of models, and the Instruct version of Jamba is coming soon to the [AI21 platform](https://www.ai21.com/studio).

model-00001-of-00021.safetensors → model-00001-of-00021.safetensors.znn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1aace34ee0da3bf95605bd150fff6d3e78110be4048a3c389b0a740354b2ccb7
-size 4951761424
+oid sha256:5672600d8bd2ea86bae0221aeebc20529037277e37962ad8e08740ee1feb8511
+size 3284450972

model-00002-of-00021.safetensors → model-00002-of-00021.safetensors.znn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0ba1de67a86329431f14f7ffa165d84055d32ce57a6d2314e3b2464eac3732dc
-size 4884669624
+oid sha256:e0dd5e53d72994e77da0a3f1f1a3f10aacdd8b174f9e4cb4da213af8a17765d7
+size 3237752699

model-00003-of-00021.safetensors → model-00003-of-00021.safetensors.znn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1abc4f16865fb78241c9453292ee3b2ca2c1e2d54ee945631da625834b95c9b2
-size 4992557120
+oid sha256:e1d36f63b738cce501b21d26e5134ac01264d2f4f92db6e9b13159ac5a8491b4
+size 3310042993

model-00004-of-00021.safetensors → model-00004-of-00021.safetensors.znn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:45fab97739a58e924791572ea3d06f9c90b9ff2a299460aaa4bd87c6e9d424f3
-size 4958853560
+oid sha256:0dcd88e533ae4d28f9e5aad9e72aa11e0f01dbf107838088bf8148d18d85749b
+size 3286974369

model-00005-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c4b0ec6e8f33e6d7b1f837cd4c25818487dcc7e478734606da28110507e51c97
-size 4975763832

model-00005-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ea2482e5fef7f57f17c328470f0351e233e6c9dc06b47edf5ce2e043013621f
+size 3298442475

model-00006-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ed98d5c3c8d7ab7352944bea09b0d54d98066cf567ba3d069da12c05575d56ed
-size 4884669616

model-00006-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87f0ce7030dcb77737feefe531c0083664ed5cbc088e08104b95f4d867028abf
+size 3237963769

model-00007-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:735be2bc568711bf42a4caebcda8288dd300b31b48fa098b00df3cf1a98e10e2
-size 4884669640

model-00007-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50e003b6df8e3656af26a0801b7606a8ecc821fcbf948fe8e51ee5c23a78e480
+size 3237782364

model-00008-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d0c8d817b2b47661d361e8b520128b3194185f756cc2204a95d642e24895ee51
-size 4992557176

model-00008-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b7474f52e0f64f4fb8eff4eb0cfaca77fa0ff9fc326d9f807ab26ac41343b33
+size 3309709335

model-00009-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e50222cf865ca5678d22574b131294303c46b249478cf70113c701f70331e999
-size 4932507176

model-00009-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7473259c0c236438c98ef334d253d08b80535d011e521a81bfb139b2fbaba49
+size 3268999814

model-00010-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b1b4b69b24ae55827b6c8b1e4a10807aa3525bc85f4d34dc002ac7440757fbf4
-size 4884669672

model-00010-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fcb191dd57edb3cb4dd26edef689b47c4585dc9821beeb33f35f4aa222c5a5f
+size 3237672285

model-00011-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:60213cac13b92ed34b93ce48e670434f22e3bf8b2b8df20c60b7bf8a9515c35c
-size 4884669696

model-00011-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:726bb3b6034717691b4c5009400e84bcdd94afb9ba2fe11444b80844621b64bb
+size 3237922482

model-00012-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:05805eacd3bb40cc9da802350409f1cb078e8b276da7e06c7a8a5ca5b26cc887
-size 4884669688

model-00012-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e309bad973392920171a6bc8ff8504d0be97a91c6ff036da65c584c5b7e56779
+size 3238025909

model-00013-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:201df979a1b34ced6cdbb7a790163412636779f1119e3845a704c489181d03d2
-size 4932507176

model-00013-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47397c014f6d5d7606818342b17501ff6f8f39083950edc8231c38bcb9d4406d
+size 3269900475

model-00014-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d0a7eb42a9ea3a385442c2e758dd5efd5dc5b913f1d10bfd37792cc963a33c93
-size 4992557152

model-00014-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b92a3a5a36d454f05bce1bee02e4d4fb1aba7524c90ab24d685f6acd4fe4f7f
+size 3309600217

model-00015-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a4b9afe4398000c28b36e3aa40c87086af673d4f8a64bfc5767941ab2008bcc9
-size 4884669688

model-00015-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:128b215d75d3d083659a1c1be9af377ff9d617abf3aceba112d77f69a1717678
+size 3237503014

model-00016-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dd1ac6cc861971c43bdf0c9c6d4c9fe72d33e5227e054a621e2e68f001419763
-size 4884669688

model-00016-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0942f4313f079e387190a07de9ca682ad2815f4060f26a2a8f14eec61abd1f24
+size 3237635876

model-00017-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:52d9eea696dd29ef413d617bbcb62a9f159e8fe8170d36e018932cef45ee281d
-size 4908522856

model-00017-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2224a44f37673e2f03e88937359d1c26895dc5ebffc20b1ce9ab617b090393c2
+size 3253266376

model-00018-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:77acada7c098e81280645ea0a9dbfa00196dca6da8946498b9907e9e376fb42d
-size 4908654000

model-00018-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ded434067927484c6ee52fb6f7e677f3d57a4c7f1ead0f34b11a18968bddaa8
+size 3253812738

model-00019-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:09e10dfd6c6459cd3460b1d667639717d3657274c1694c19a6fdbac1be6a76bf
-size 4992557168

model-00019-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3382ceb4b2101b8714f0b1a41e06debfdaa947dd8a1ed2257314c837a7b1bbf9
+size 3308468840

model-00020-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2bd5c27b2cca6e06f7b4497ce8c9b1522a64846817a871bad274d08507960ed0
-size 4884669696

model-00020-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7ce533fe923573760a664973927816bbf39efa9c39a7cb10d0539068873869c
+size 3237872986

model-00021-of-00021.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a47ef23db8deb5364da676a40dc3dcb011fb9d9ceef13ba044c176e9a83ac1e3
-size 4647318576

model-00021-of-00021.safetensors.znn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:322bce95bcb7f0fabc5fd1f610a4b2aca6eed1bf576104e01aebb17b1337f2a9
+size 3080471330

zipnn_compress_file.py

@@ -0,0 +1,154 @@
+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
+
+
+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,
+):
+    import zipnn
+
+    streaming_chunk_size = parse_streaming_chunk_size(streaming_chunk_size)
+    full_path = input_file
+    if not os.path.exists(full_path):
+        print("File not found")
+        return
+    if delete:
+        print(f"Deleting {full_path}...")
+        os.remove(full_path)
+    else:
+        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"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}"
+        )
+
+
+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.",
+    )
+    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
+
+    check_and_install_zipnn()
+    compress_file(args.input_file, **optional_kwargs)

zipnn_compress_path.py

@@ -0,0 +1,253 @@
+import os
+import subprocess
+import sys
+import argparse
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    as_completed,
+)
+from zipnn_compress_file import compress_file
+import zipnn
+
+sys.path.append(
+    os.path.abspath(
+        os.path.join(
+            os.path.dirname(__file__), ".."
+        )
+    )
+)
+
+
+KB = 1024
+MB = 1024 * 1024
+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",
+                "--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_files_with_suffix(
+    suffix,
+    dtype="",
+    streaming_chunk_size=1048576,
+    path=".",
+    delete=False,
+    r=False,
+    force=False,
+    max_processes=1,
+):
+
+    file_list = []
+    streaming_chunk_size = (
+        parse_streaming_chunk_size(
+            streaming_chunk_size
+        )
+    )
+    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
+                ):
+                    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)
+
+    with ProcessPoolExecutor(
+        max_workers=max_processes
+    ) as executor:
+        future_to_file = {
+            executor.submit(
+                compress_file,
+                file,
+                dtype,
+                streaming_chunk_size,
+                delete,
+                True,
+            ): 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"File {file} generated an exception: {exc}"
+                    )
+
+                if file_list:
+                    next_file = file_list.pop(0)
+                    future_to_file[
+                        executor.submit(
+                            compress_file,
+                            next_file,
+                            dtype,
+                            streaming_chunk_size,
+                            delete,
+                            True,
+                        )
+                    ] = next_file
+
+    if not files_found:
+        print(
+            f"No files with the suffix '{suffix}' found."
+        )
+
+
+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.",
+    )
+    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
+        )
+
+    check_and_install_zipnn()
+    compress_files_with_suffix(
+        args.suffix, **optional_kwargs
+    )

zipnn_decompress_file.py

@@ -0,0 +1,84 @@
+import os
+import subprocess
+import sys
+import argparse
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+
+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, dtype="", delete=False, force=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):
+        if delete:
+            print(f"Deleting {input_file}...")
+            os.remove(input_file)
+        else:
+            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]
+
+            if dtype:
+                zpn = zipnn.ZipNN(is_streaming=True, bytearray_dtype="float32")
+            else:
+                zpn = zipnn.ZipNN(is_streaming=True)
+
+            with open(input_file, "rb") as infile, open(output_file, "wb") as outfile:
+                d_data = b""
+                chunk = infile.read()
+                d_data += zpn.decompress(chunk)
+                outfile.write(d_data)
+                print(f"Decompressed {input_file} to {output_file}")
+
+    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(
+        "--float32", action="store_true", help="A flag that triggers float32 compression."
+    )
+    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."
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    if args.delete:
+        optional_kwargs["delete"] = args.delete
+    if args.force:
+        optional_kwargs["force"] = args.force
+
+    decompress_file(args.input_file, **optional_kwargs)

zipnn_decompress_path.py

@@ -0,0 +1,184 @@
+import os
+import sys
+import argparse
+import subprocess
+import zipnn
+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__),
+            "..",
+        )
+    )
+)
+
+
+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_zpn_files(
+    dtype="",
+    path=".",
+    delete=False,
+    force=False,
+    max_processes=1,
+):
+
+    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
+                ):
+                    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)
+
+    with ProcessPoolExecutor(
+        max_workers=max_processes
+    ) as executor:
+        for file in file_list[:max_processes]:
+            future_to_file = {
+                executor.submit(
+                    decompress_file,
+                    file,
+                    dtype,
+                    delete,
+                    True,
+                ): 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"File {file} generated an exception: {exc}"
+                        )
+
+                    if file_list:
+                        next_file = file_list.pop(
+                            0
+                        )
+                        future_to_file[
+                            executor.submit(
+                                decompress_file,
+                                next_file,
+                                dtype,
+                                delete,
+                                True,
+                            )
+                        ] = next_file
+                        #
+
+
+if __name__ == "__main__":
+    check_and_install_zipnn()
+
+    parser = argparse.ArgumentParser(
+        description="Compresses all .znn files. (optional) dtype."
+    )
+    parser.add_argument(
+        "--float32",
+        action="store_true",
+        help="A flag that triggers float32 compression.",
+    )
+    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.",
+    )
+    args = parser.parse_args()
+    optional_kwargs = {}
+    if args.float32:
+        optional_kwargs["dtype"] = 32
+    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
+        )
+
+    decompress_zpn_files(**optional_kwargs)