compressed-tensors 0.3.0__py3-none-any.whl

compressed_tensors/utils/safetensors_load.py

@@ -0,0 +1,237 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import re
+import struct
+from typing import Dict, List, Optional
+
+from safetensors import safe_open
+from torch import Tensor
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, cached_file
+
+
+__all__ = [
+    "get_safetensors_folder",
+    "get_safetensors_header",
+    "match_param_name",
+    "merge_names",
+    "get_weight_mappings",
+    "get_nested_weight_mappings",
+    "get_quantization_state_dict",
+]
+
+
+def get_safetensors_folder(
+    pretrained_model_name_or_path: str, cache_dir: Optional[str] = None
+) -> str:
+    """
+    Given a Hugging Face stub or a local path, return the folder containing the
+    safetensors weight files
+
+    :param pretrained_model_name_or_path: local path to model or HF stub
+    :param cache_dir: optional cache dir to search through, if none is specified the
+    model will be searched for in the default TRANSFORMERS_CACHE
+    :return: local folder containing model data
+    """
+    if os.path.exists(pretrained_model_name_or_path):
+        # argument is a path to a local folder
+        return os.path.abspath(pretrained_model_name_or_path)
+
+    safetensors_path = cached_file(
+        pretrained_model_name_or_path,
+        SAFE_WEIGHTS_NAME,
+        cache_dir=cache_dir,
+        _raise_exceptions_for_missing_entries=False,
+    )
+    index_path = cached_file(
+        pretrained_model_name_or_path,
+        SAFE_WEIGHTS_INDEX_NAME,
+        cache_dir=cache_dir,
+        _raise_exceptions_for_missing_entries=False,
+    )
+    if safetensors_path is not None:
+        # found a single cached safetensors file
+        return os.path.split(safetensors_path)[0]
+    if index_path is not None:
+        # found a cached safetensors weight index file
+        return os.path.split(index_path)[0]
+
+    # model weights could not be found locally or cached from HF Hub
+    raise ValueError(
+        "Could not locate safetensors weight or index file from "
+        f"{pretrained_model_name_or_path}."
+    )
+
+
+def get_safetensors_header(safetensors_path: str) -> Dict[str, str]:
+    """
+    Extracts the metadata from a safetensors file as JSON
+
+    :param safetensors_path: path to a safetensors file
+    :return: dictionary of metadata extracted from the safetensors file
+    """
+    with open(safetensors_path, "rb") as f:
+        length_of_header = struct.unpack("<Q", f.read(8))[0]
+        header_data = f.read(length_of_header)
+        header = json.loads(header_data)
+
+    return header
+
+
+def match_param_name(full_name: str, param_name: str) -> str:
+    """
+    Helper function extracting the uncompressed parameterized layer name from a
+    compressed name. Assumes the compressed name was merged using merge_names.
+
+    :param full_name: full name of parameter in compressed model
+    :param param_name: compression paramater name
+    :return: uncompressed name of the uncompressed parameterized layer
+    """
+    pattern = r"^(.*)\." + param_name + r"$"
+    regex = re.findall(pattern, full_name)
+    if len(regex) == 0:
+        return None
+    return regex[0]
+
+
+def merge_names(parent_name: str, child_name: str) -> str:
+    """
+    Helper function for merging an uncompressed parameterized layer name with a
+    compression parameter. Names merged with this function can then be parsed by
+    match_param_name.
+
+    :param parent_name: uncompressed parameterized layer name
+    :param child_name: compression parameter name
+    :return: merged compressed name
+    """
+    return parent_name + "." + child_name
+
+
+def get_weight_mappings(path_to_model_or_tensors: str) -> Dict[str, str]:
+    """
+    Takes a path to a state dict saved in safetensors format and returns a mapping
+    from parameterized layer name to file location.
+
+    {
+        layer.weight.bitmask: file_location,
+        layer.weight.row_offsets: file_location,
+        layer.weight.shape: file_location,
+        layer.weight.compressed: file_location
+    }
+
+    This generalizes to cases where the model is split into multiple safetensors files
+
+    :param path_to_model_or_tensors: path to directory that contains
+        safetensors (must contain either a single file or multiple files with an index),
+        or a path to a single safetensors file
+    :return: mapping of parameterized layer name to file location
+    """
+
+    if os.path.isfile(path_to_model_or_tensors):
+        # we have a single safetensors file to read
+        header = get_safetensors_header(path_to_model_or_tensors)
+        for key in header.keys():
+            header[key] = path_to_model_or_tensors
+        header.pop("__metadata__", None)
+    else:
+        # we have a directory with multiple safetensors files
+        safetensors_path = os.path.join(path_to_model_or_tensors, SAFE_WEIGHTS_NAME)
+        index_path = os.path.join(path_to_model_or_tensors, SAFE_WEIGHTS_INDEX_NAME)
+        if os.path.exists(safetensors_path):
+            # we have a single safetensors file to read
+            header = get_safetensors_header(safetensors_path)
+            for key in header.keys():
+                header[key] = SAFE_WEIGHTS_NAME
+            header.pop("__metadata__", None)
+        elif os.path.exists(index_path):
+            # we have multiple safetensors file, read from index
+            with open(index_path, "r", encoding="utf-8") as f:
+                index = json.load(f)
+            header = index["weight_map"]
+        else:
+            raise ValueError(
+                "Could not find a safetensors weight "
+                f"or index file at {path_to_model_or_tensors}"
+            )
+
+        # convert weight locations to full paths
+        for key, value in header.items():
+            header[key] = os.path.join(path_to_model_or_tensors, value)
+
+    return header
+
+
+def get_nested_weight_mappings(
+    model_path: str, params_to_nest: List[str]
+) -> Dict[str, Dict[str, str]]:
+    """
+    Takes a path to a state dict saved in safetensors format and returns a nested
+    mapping from uncompressed parameterized layer names to the file locations of each
+    of the layers compression parameters.
+
+    layer.weight: {
+        bitmask: file_location,
+        row_offsets: file_location,
+        shape: file_location,
+        compressed: file_location
+    }
+
+    This generalizes to cases where the model is split into multiple safetensors files
+
+    :param model_path: path to safetensors state dict, must contain either a single
+    safetensors file or multiple files with an index
+    :return: nested mapping of parameterized layer name to file location
+    """
+    weight_mappings = get_weight_mappings(model_path)
+
+    nested_weight_mappings = {}
+    for key in weight_mappings.keys():
+        for param_name in params_to_nest:
+            maybe_match = match_param_name(key, param_name)
+            if maybe_match is not None:
+                dense_param = maybe_match
+                if dense_param not in nested_weight_mappings:
+                    nested_weight_mappings[dense_param] = {}
+                nested_weight_mappings[dense_param][param_name] = weight_mappings[key]
+
+    return nested_weight_mappings
+
+
+def get_quantization_state_dict(model_path: str) -> Dict[str, Tensor]:
+    weight_mappings = get_weight_mappings(model_path)
+    state_dict = {}
+    for weight_name, safe_path in weight_mappings.items():
+        if not _is_quantization_weight(weight_name):
+            continue
+        with safe_open(safe_path, framework="pt", device="cpu") as f:
+            state_dict[weight_name] = f.get_tensor(weight_name)
+
+    return state_dict
+
+
+def _is_quantization_weight(name: str) -> bool:
+    """
+    Checks is a parameter name is associated with a quantization parameter
+
+    :param name: parameter name to check
+    :return: True if parameter name is a quantization parameter, else False
+    """
+    if name.endswith("_scale"):
+        return True
+    if name.endswith("zero_point"):
+        return True
+
+    return False

compressed_tensors-0.3.0.dist-info/METADATA

@@ -0,0 +1,22 @@
+Metadata-Version: 2.1
+Name: compressed-tensors
+Version: 0.3.0
+Summary: Library for utilization of compressed safetensors of neural network models
+Home-page: UNKNOWN
+Author: Neuralmagic, Inc.
+Author-email: support@neuralmagic.com
+License: UNKNOWN
+Platform: UNKNOWN
+Requires-Dist: pydantic <2.7
+Requires-Dist: torch >=1.7.0
+Requires-Dist: transformers <=4.40
+Provides-Extra: dev
+Requires-Dist: black ==22.12.0 ; extra == 'dev'
+Requires-Dist: flake8 >=3.8.3 ; extra == 'dev'
+Requires-Dist: isort ==5.8.0 ; extra == 'dev'
+Requires-Dist: nbconvert >=7.16.3 ; extra == 'dev'
+Requires-Dist: pytest >=6.0.0 ; extra == 'dev'
+Requires-Dist: wheel >=0.36.2 ; extra == 'dev'
+
+UNKNOWN
+

compressed_tensors-0.3.0.dist-info/RECORD

@@ -0,0 +1,37 @@
+compressed_tensors/__init__.py,sha256=SV1csvHUVCd8kHXz6UDZim1HZ_fAVG3vfk-j_4Bb6hY,789
+compressed_tensors/base.py,sha256=8zbgK87LpHkKoSknM55svXCT4E4dLLjPijwF9HfzmsQ,717
+compressed_tensors/compressors/__init__.py,sha256=3ZHKWSIWTjMx8XXgLtoP9JaVaCTvRecguLZTxLAAkKk,898
+compressed_tensors/compressors/base.py,sha256=F1smyJ6x2Sfq43tuP0QE9wZuhVqnewq-XUFPMtdU9yQ,2936
+compressed_tensors/compressors/dense.py,sha256=_VTusI3XjaY-zOdB_d7z4zOgPTJi9TJZZHF13g9ulS4,1263
+compressed_tensors/compressors/helpers.py,sha256=kSseqbwnu3JHZUKH8u4kQo5bmd87FvCcmWe0u2ikysA,6421
+compressed_tensors/compressors/sparse_bitmask.py,sha256=PYAK_Hcy2T57zlbpwl1FYkslluIr2x-d0Rh048YAtpI,8639
+compressed_tensors/config/__init__.py,sha256=ZBqWn3r6ku1qfmlHHYp0mQueY0i7Pwhr9rbQk9dDlMc,704
+compressed_tensors/config/base.py,sha256=IP-3Y416w-811WozDzKHycIBXjdlG4Ddy7vpbwhOPD8,1373
+compressed_tensors/config/dense.py,sha256=xtkri7DkP7USu44FnSoTgTSqdGegCBtjRf3DfblSEL0,1311
+compressed_tensors/config/sparse_bitmask.py,sha256=y8fmQaOoGjIiI4FR6BJjfIqisAcqNQ_zjKyjT75bXwY,1284
+compressed_tensors/quantization/__init__.py,sha256=83J5bPB7PavN2TfCoW7_vEDhfYpm4TDrqYO9vdSQ5bk,760
+compressed_tensors/quantization/quant_args.py,sha256=dxSrq0_88ORQXcyIMYqoMZJvYEjnqdYl37f7lgZQqhw,2742
+compressed_tensors/quantization/quant_config.py,sha256=DWx8ae3gDlw99zAn3MUN9I4qeksbbmITmOXHRynqPB8,6650
+compressed_tensors/quantization/quant_scheme.py,sha256=X3oqmZPiIKtX5tEKKUj-0N6hB68NeiU2b1GcQEQPadQ,1480
+compressed_tensors/quantization/lifecycle/__init__.py,sha256=fM9XBtPgJX6z54PTm3Sd0SpK5od95ibwaSf2FFR8DqE,772
+compressed_tensors/quantization/lifecycle/apply.py,sha256=WXUL3q1g0s244k0wuqGYZPXTXiscdyrp7RScN2j_KGA,6651
+compressed_tensors/quantization/lifecycle/calibration.py,sha256=mLns4jlaWmBwOW8Jtlm5bMX-JET1AiZYUBO7qa-XuxI,1776
+compressed_tensors/quantization/lifecycle/forward.py,sha256=hnjk7pocZLDhLdMx237FKayYdvsdKbYSjTmSN5xbQO8,4599
+compressed_tensors/quantization/lifecycle/frozen.py,sha256=NHNmlDIaxurifqeI_qZC8xa4BstQsBNdOCXJjRzAfNU,1596
+compressed_tensors/quantization/lifecycle/initialize.py,sha256=8pifqZQSgVqWYI_Qtv6QfBICPbCTFHy48OWPeQsxEHQ,3578
+compressed_tensors/quantization/observers/__init__.py,sha256=DNH31NQYrIBBcmHsMyFA6whh4pbRsLwuNa6L8AeXaGc,745
+compressed_tensors/quantization/observers/base.py,sha256=O76dAxkin7bB602e9kjmxc84p71-PxBtjIq5L69xplI,2786
+compressed_tensors/quantization/observers/helpers.py,sha256=SxvOf9zwZ9NDRC3E4Xm7z3RqHcbcPtCABLKX9GnGGHM,2109
+compressed_tensors/quantization/observers/memoryless.py,sha256=3f6bUlcf5mzOHPkTRhoQ7Zd8xu_pUmj8e3Y85fGysSU,1848
+compressed_tensors/quantization/observers/min_max.py,sha256=uAcZd5aY6WKM-KumTb2ybX28s8iKGVy6Nrje5Sddqew,2439
+compressed_tensors/quantization/utils/__init__.py,sha256=VdtEmP0bvuND_IGQnyqUPc5lnFp-1_yD7StKSX4x80w,656
+compressed_tensors/quantization/utils/helpers.py,sha256=N_wYfrPcFr__Q1mn6mHoNUTclwpTW8P5PDHkR7GvXWo,3694
+compressed_tensors/registry/__init__.py,sha256=FwLSNYqfIrb5JD_6OK_MT4_svvKTN_nEhpgQlQvGbjI,658
+compressed_tensors/registry/registry.py,sha256=fxjOjh2wklCvJhQxwofdy-zV8q7MkQ85SLG77nml2iA,11890
+compressed_tensors/utils/__init__.py,sha256=5DrYjoZbaEvSkJcC-GRSbM_RBHVF4tG9gMd3zsJnjLw,665
+compressed_tensors/utils/helpers.py,sha256=wLgiPrk7Vn29AijOGQGk3UnXItRd1jpROS6FxHoC4VQ,5530
+compressed_tensors/utils/safetensors_load.py,sha256=wo9UirGrGlenBqZeqotvpCT7D5MEdjCo2J3HeRaIFoU,8502
+compressed_tensors-0.3.0.dist-info/METADATA,sha256=REuv9Fm1mYXww7d8VOWwPWYkAhImzn_PEhiPn_V5Hdo,673
+compressed_tensors-0.3.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+compressed_tensors-0.3.0.dist-info/top_level.txt,sha256=w2i-GyPs2s1UwVxvutSvN_lM22SXC2hQFBmoMcPnV7Y,19
+compressed_tensors-0.3.0.dist-info/RECORD,,

compressed_tensors-0.3.0.dist-info/WHEEL

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

compressed_tensors-0.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+compressed_tensors