initial commit

README.md

@@ -80,3 +80,5 @@ save_compressed_model(model, "compressed_model.safetensors", compression_format=
 # load compressed model weights (`dict` turns generator into a dictionary)
 state_dict = dict(load_compressed("compressed_model.safetensors", compression_config))
 ```
+
+For more in-depth tutorial on bitmask compression, refer to the notebook: `examples/bitmask_compression.ipynb`

setup.py

@@ -32,7 +32,7 @@ def _setup_extras() -> Dict:
 
 setup(
     name="compressed-tensors",
-    version="0.3.1",
+    version="0.3.2",
     author="Neuralmagic, Inc.",
     author_email="support@neuralmagic.com",
     license="Apache 2.0",

src/compressed_tensors/compressors/dense.py

@@ -29,6 +29,6 @@ def compress(self, model_state: Dict[str, Tensor]) -> Dict[str, Tensor]:
         return model_state
 
     def decompress(
-        self, path_to_model_or_tensors: str, device: str
+        self, path_to_model_or_tensors: str, device: str = "cpu"
     ) -> Generator[Tuple[str, Tensor], None, None]:
         return iter([])

src/compressed_tensors/compressors/sparse_bitmask.py

@@ -75,8 +75,9 @@ def decompress(
         self, path_to_model_or_tensors: str, device: str = "cpu"
     ) -> Generator[Tuple[str, Tensor], None, None]:
         """
-        Reads a bitmask compressed state dict located at path_to_model_or_tensors
-        and returns a generator for sequentially decompressing back to a dense state dict
+        Reads a bitmask compressed state dict located
+        at path_to_model_or_tensors and returns a generator
+        for sequentially decompressing back to a dense state dict
 
         :param model_path: path to compressed safetensors model (directory with
             one or more safetensors files) or compressed tensors file