compressed-tensors-nightly 0.9.1.20250204__py3-none-any.whl → 0.9.1.20250206__py3-none-any.whl

compressed_tensors/compressors/base.py

@@ -77,6 +77,15 @@ class BaseCompressor(RegistryMixin, ABC):
         """
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        raise NotImplementedError()
+
     @abstractmethod
     def compress(
         self,

compressed_tensors/compressors/quantized_compressors/base.py

@@ -144,7 +144,7 @@ class BaseQuantizationCompressor(BaseCompressor):
 
     def _decompress_from_path(self, path_to_model, names_to_scheme, device):
         weight_mappings = get_nested_weight_mappings(
-            path_to_model, self.COMPRESSION_PARAM_NAMES
+            path_to_model, self.compression_param_names
         )
         for weight_name in weight_mappings.keys():
             weight_data = {}
@@ -161,7 +161,7 @@ class BaseQuantizationCompressor(BaseCompressor):
 
     def _decompress_from_state_dict(self, state_dict, names_to_scheme):
         weight_mappings = get_nested_mappings_from_state_dict(
-            state_dict, self.COMPRESSION_PARAM_NAMES
+            state_dict, self.compression_param_names
         )
         for weight_name in weight_mappings.keys():
             weight_data = {}

compressed_tensors/compressors/quantized_compressors/naive_quantized.py

@@ -41,12 +41,18 @@ class NaiveQuantizationCompressor(BaseQuantizationCompressor):
     type to the type specified by the layer's QuantizationArgs.
     """
 
-    COMPRESSION_PARAM_NAMES = [
-        "weight",
-        "weight_scale",
-        "weight_zero_point",
-        "weight_g_idx",
-    ]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return (
+            "weight",
+            "weight_scale",
+            "weight_zero_point",
+            "weight_g_idx",
+        )
 
     def compression_param_info(
         self,

compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -36,13 +36,19 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
     Compresses a quantized model by packing every eight 4-bit weights into an int32
     """
 
-    COMPRESSION_PARAM_NAMES = [
-        "weight_packed",
-        "weight_scale",
-        "weight_zero_point",
-        "weight_g_idx",
-        "weight_shape",
-    ]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return (
+            "weight_packed",
+            "weight_scale",
+            "weight_zero_point",
+            "weight_g_idx",
+            "weight_shape",
+        )
 
     def compression_param_info(
         self,

compressed_tensors/compressors/sparse_compressors/base.py

@@ -30,8 +30,7 @@ _LOGGER: logging.Logger = logging.getLogger(__name__)
 class BaseSparseCompressor(BaseCompressor):
     """
     Base class representing a sparse compression algorithm. Each child class should
-    implement compression_param_info, compress_weight and decompress_weight; child
-    classes should also define COMPRESSION_PARAM_NAMES.
+    implement compression_param_names, compress_weight and decompress_weight;
 
     Compressors support compressing/decompressing a full module state dict or a single
     quantized PyTorch leaf module.
@@ -113,7 +112,7 @@ class BaseSparseCompressor(BaseCompressor):
         """
         weight_mappings, ignored_params = get_nested_weight_mappings(
             path_to_model_or_tensors,
-            self.COMPRESSION_PARAM_NAMES,
+            self.compression_param_names,
             return_unmatched_params=True,
         )
         for weight_name in weight_mappings.keys():

compressed_tensors/compressors/sparse_compressors/dense.py

@@ -25,6 +25,14 @@ class DenseCompressor(BaseCompressor):
     Identity compressor for dense models, returns the original state_dict
     """
 
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return ()
+
     def compress(self, model_state: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor]:
         return model_state
 

compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

@@ -40,11 +40,17 @@ class Sparse24BitMaskCompressor(BaseSparseCompressor):
     values tensor, with their locations stored in a 2d bitmask
     """
 
-    COMPRESSION_PARAM_NAMES = [
-        "shape",
-        "compressed",
-        "bitmask",
-    ]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return (
+            "shape",
+            "compressed",
+            "bitmask",
+        )
 
     def compress_weight(self, name, value):
         bitmask_tensor = Sparse24BitMaskTensor.from_dense(

compressed_tensors/compressors/sparse_compressors/sparse_bitmask.py

@@ -38,7 +38,13 @@ class BitmaskCompressor(BaseSparseCompressor):
     values tensor, with their locations stored in a 2d bitmask
     """
 
-    COMPRESSION_PARAM_NAMES = ["shape", "compressed", "bitmask", "row_offsets"]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return ("shape", "compressed", "bitmask", "row_offsets")
 
     def compress_weight(self, name, value):
         bitmask_tensor = BitmaskTensor.from_dense(value)

compressed_tensors/compressors/sparse_quantized_compressors/marlin_24.py

@@ -42,8 +42,6 @@ class Marlin24Compressor(BaseCompressor):
     Marlin24 kernel. Decompression is not implemented for this compressor.
     """
 
-    COMPRESSION_PARAM_NAMES = ["weight_packed", "scale_packed", "meta"]
-
     @staticmethod
     def validate_quant_compatability(
         model_quant_args: Dict[str, QuantizationArgs]
@@ -105,6 +103,14 @@ class Marlin24Compressor(BaseCompressor):
 
         return True
 
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return ("weight_packed", "scale_packed", "meta")
+
     def compress(
         self,
         model_state: Dict[str, Tensor],

compressed_tensors/utils/safetensors_load.py

@@ -16,7 +16,7 @@ import json
 import os
 import re
 import struct
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, Iterable, Optional, Tuple, Union
 
 from safetensors import safe_open
 from torch import Tensor
@@ -180,7 +180,9 @@ def get_weight_mappings(path_to_model_or_tensors: str) -> Dict[str, str]:
 
 
 def get_nested_weight_mappings(
-    model_path: str, params_to_nest: List[str], return_unmatched_params: bool = False
+    model_path: str,
+    params_to_nest: Iterable[str],
+    return_unmatched_params: bool = False,
 ) -> Union[NestedWeightMappingType, Tuple[NestedWeightMappingType, WeightMappingType]]:
     """
     Takes a path to a state dict saved in safetensors format and returns a nested
@@ -211,7 +213,7 @@ def get_nested_weight_mappings(
 
     :param model_path: Path to the safetensors state dict, must contain either a
         single safetensors file or multiple files with an index.
-    :param params_to_nest: List of parameter names to nest.
+    :param params_to_nest: Iterable of parameter names to nest.
     :param return_unmatched_params: If True, return a second dictionary containing
         the remaining parameters that were not matched to the params_to_nest.
     :return:
@@ -247,7 +249,7 @@ def get_nested_weight_mappings(
 
 
 def get_nested_mappings_from_state_dict(
-    state_dict, params_to_nest
+    state_dict, params_to_nest: Iterable[str]
 ) -> NestedWeightMappingType:
     """
     Takes a state dict and returns a nested mapping from uncompressed
@@ -262,7 +264,7 @@ def get_nested_mappings_from_state_dict(
     }
 
     :param state_dict: state dict of the model
-    :param params_to_nest: List of parameter names to nest.
+    :param params_to_nest: Iterable of parameter names to nest.
     :return: Nested mapping of parameterized layer names to the value of
         each layer's compression parameters.
     """

{compressed_tensors_nightly-0.9.1.20250204.dist-info → compressed_tensors_nightly-0.9.1.20250206.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: compressed-tensors-nightly
-Version: 0.9.1.20250204
+Version: 0.9.1.20250206
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/neuralmagic/compressed-tensors
 Author: Neuralmagic, Inc.
@@ -11,8 +11,6 @@ License-File: LICENSE
 Requires-Dist: torch>=1.7.0
 Requires-Dist: transformers
 Requires-Dist: pydantic>=2.0
-Provides-Extra: accelerate
-Requires-Dist: accelerate; extra == "accelerate"
 Provides-Extra: dev
 Requires-Dist: black==22.12.0; extra == "dev"
 Requires-Dist: isort==5.8.0; extra == "dev"
@@ -20,6 +18,17 @@ Requires-Dist: wheel>=0.36.2; extra == "dev"
 Requires-Dist: flake8>=3.8.3; extra == "dev"
 Requires-Dist: pytest>=6.0.0; extra == "dev"
 Requires-Dist: nbconvert>=7.16.3; extra == "dev"
+Provides-Extra: accelerate
+Requires-Dist: accelerate; extra == "accelerate"
+Dynamic: author
+Dynamic: author-email
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: summary
 
 # compressed-tensors
 

{compressed_tensors_nightly-0.9.1.20250204.dist-info → compressed_tensors_nightly-0.9.1.20250206.dist-info}/RECORD

@@ -2,21 +2,21 @@ compressed_tensors/__init__.py,sha256=UtKmifNeBCSE2TZSAfduVNNzHY-3V7bLjZ7n7RuXLO
 compressed_tensors/base.py,sha256=73HYH7HY7O2roC89yG_piPFnZwrBfn_i7HmKl90SKc0,875
 compressed_tensors/version.py,sha256=4fhjdvtnsinY2nzy2jDu0CL4Tytgk6EVPNsww7-spZU,1586
 compressed_tensors/compressors/__init__.py,sha256=smSygTSfcfuujRrAXDc6uZm4L_ccV1tWZewqVnOb4lM,825
-compressed_tensors/compressors/base.py,sha256=D9TNwQcjanDiAHODPbg8JUqc66e3j50rctY7A708NEs,6743
+compressed_tensors/compressors/base.py,sha256=x8dQrWVEurynXw03yHJZTaAmrRTOsdZJoHjmvs0IKwk,7002
 compressed_tensors/compressors/helpers.py,sha256=OK6qxX9j3bHwF9JfIYSGMgBJe2PWjlTA3byXKCJaTIQ,5431
 compressed_tensors/compressors/model_compressors/__init__.py,sha256=5RGGPFu4YqEt_aOdFSQYFYFDjcZFJN0CsMqRtDZz3Js,666
 compressed_tensors/compressors/model_compressors/model_compressor.py,sha256=3WyzAW2Rm_uLprxwO2QH6FR76W6Mk4r2yedayaSZHhw,18396
 compressed_tensors/compressors/quantized_compressors/__init__.py,sha256=09UJq68Pht6Bf-4iP9xYl3tetKsncNPHD8IAGbePsr4,714
-compressed_tensors/compressors/quantized_compressors/base.py,sha256=LVqSSqSjGi8LB-X13zC_0AFHc8BobGQVC0zjInDhOWE,7217
-compressed_tensors/compressors/quantized_compressors/naive_quantized.py,sha256=fahmPJFz49rVS7q705uQwZ0kUtdP46GuXR7nPr6uIqI,4943
-compressed_tensors/compressors/quantized_compressors/pack_quantized.py,sha256=hMGGQhf1fTNTE0m7QJwet_dEEAWzuT1caiVieZ1KqnM,8298
+compressed_tensors/compressors/quantized_compressors/base.py,sha256=cp8S1Kr3HhlMHIz7k4vGo-qxxdknEC3qP1QLIhNnwRA,7217
+compressed_tensors/compressors/quantized_compressors/naive_quantized.py,sha256=fd0KlkSx6bvZ3xwIkK3jEUdPSUPs56Eua4dEDOtzKW0,5150
+compressed_tensors/compressors/quantized_compressors/pack_quantized.py,sha256=zH2PocRe_T5yt1-3kLdZH9AUQWQyaVOi4U9nEJiYaWA,8509
 compressed_tensors/compressors/sparse_compressors/__init__.py,sha256=Atuz-OdEgn8OCUhx7Ovd6gXdyImAI186uCR-uR0t_Nk,737
-compressed_tensors/compressors/sparse_compressors/base.py,sha256=9e841MQWr0j8m33ejDw_jP5_BIpQ5099x9_pvuZ-Nr0,5944
-compressed_tensors/compressors/sparse_compressors/dense.py,sha256=lSKNWRx6H7aUqaJj1j4qbXk8Gkm1UohbnvW1Rvq6Ra4,1284
-compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py,sha256=_g139pe4iAFn5jvGIEk4v-qMoyx9ID6E88vriPSNYV4,8604
-compressed_tensors/compressors/sparse_compressors/sparse_bitmask.py,sha256=7zSr9bqkpuH1ivQpxtYBNxXIoElal7Jo1nSKpZN_IFk,5633
+compressed_tensors/compressors/sparse_compressors/base.py,sha256=CVWbs3sd7GKJEoWOIKImABQ01VOTX8dlF2AQaEVPotw,5883
+compressed_tensors/compressors/sparse_compressors/dense.py,sha256=_uW_HISeDNz4yboSZWoh6GwrkUE6HFibzPQSKrHOCkg,1505
+compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py,sha256=mEKSSgpXookqYSJw3mlyP6cYYKD-eaIvpQMvi4JO6TY,8807
+compressed_tensors/compressors/sparse_compressors/sparse_bitmask.py,sha256=S8vW0FI9ep_XtUQOxj0P5utJt3vKEYOHjWEPp-Xd9aY,5820
 compressed_tensors/compressors/sparse_quantized_compressors/__init__.py,sha256=4f_cwcKXB1nVVMoiKgTFAc8jAPjPLElo-Df_EDm1_xw,675
-compressed_tensors/compressors/sparse_quantized_compressors/marlin_24.py,sha256=BMIQWTLlnUvxy14iEJegtiP75WHJeOVojey9mKOK1hE,9427
+compressed_tensors/compressors/sparse_quantized_compressors/marlin_24.py,sha256=xY0CdHXAzVHeDeSCD_I-5UZKcntKzd3FiKSP-ZqcSBs,9614
 compressed_tensors/config/__init__.py,sha256=8sOoZ6xvYSC79mBvEtO8l6xk4PC80d29AnnJiGMrY2M,737
 compressed_tensors/config/base.py,sha256=R3iUmFf1MslEjin5LgwQbmfJHIsS7Uw0UIxfn780uqY,3479
 compressed_tensors/config/dense.py,sha256=NgSxnFCnckU9-iunxEaqiFwqgdO7YYxlWKR74jNbjks,1317
@@ -43,10 +43,10 @@ compressed_tensors/utils/helpers.py,sha256=RrNvzD08naEjEiXdU-FdZjQVda1nQywu1hA_G
 compressed_tensors/utils/offload.py,sha256=H4aAg21zUvJM2uwE6QCNYazX_p_o41yQUAgLLWBqR0w,14079
 compressed_tensors/utils/permutations_24.py,sha256=kx6fsfDHebx94zsSzhXGyCyuC9sVyah6BUUir_StT28,2530
 compressed_tensors/utils/permute.py,sha256=V6tJLKo3Syccj-viv4F7ZKZgJeCB-hl-dK8RKI_kBwI,2355
-compressed_tensors/utils/safetensors_load.py,sha256=fBuoHVPoBt1mkvqFJ60zQIASX_4nhl0-6QfFS27NY8I,11430
+compressed_tensors/utils/safetensors_load.py,sha256=5SeM2hzLh77Ne8Vk7qR6-km7cf8bhov41ExpWITqX3A,11470
 compressed_tensors/utils/semi_structured_conversions.py,sha256=XKNffPum54kPASgqKzgKvyeqWPAkair2XEQXjkp7ho8,13489
-compressed_tensors_nightly-0.9.1.20250204.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-compressed_tensors_nightly-0.9.1.20250204.dist-info/METADATA,sha256=cTiBx68KqIk79WTQy_oQESgOqtULcIQ8Mg_7Y5jba2E,6799
-compressed_tensors_nightly-0.9.1.20250204.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-compressed_tensors_nightly-0.9.1.20250204.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
-compressed_tensors_nightly-0.9.1.20250204.dist-info/RECORD,,
+compressed_tensors_nightly-0.9.1.20250206.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+compressed_tensors_nightly-0.9.1.20250206.dist-info/METADATA,sha256=h-ujgkOd3-2jlbqgYI2gITSWqfgFu8M3z8qix5RJFIs,6992
+compressed_tensors_nightly-0.9.1.20250206.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+compressed_tensors_nightly-0.9.1.20250206.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors_nightly-0.9.1.20250206.dist-info/RECORD,,

{compressed_tensors_nightly-0.9.1.20250204.dist-info → compressed_tensors_nightly-0.9.1.20250206.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.45.1)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any