Delete safetensors_util.py

safetensors_util.py

@@ -1,81 +0,0 @@
-import base64
-import pickle
-from dataclasses import dataclass
-from typing import Dict, Optional, Tuple
-
-import safetensors.torch
-import torch
-
-from .aliases import PathOrStr
-
-__all__ = [
-    "state_dict_to_safetensors_file",
-    "safetensors_file_to_state_dict",
-]
-
-
-@dataclass(eq=True, frozen=True)
-class STKey:
-    keys: Tuple
-    value_is_pickled: bool
-
-
-def encode_key(key: STKey) -> str:
-    b = pickle.dumps((key.keys, key.value_is_pickled))
-    b = base64.urlsafe_b64encode(b)
-    return str(b, "ASCII")
-
-
-def decode_key(key: str) -> STKey:
-    b = base64.urlsafe_b64decode(key)
-    keys, value_is_pickled = pickle.loads(b)
-    return STKey(keys, value_is_pickled)
-
-
-def flatten_dict(d: Dict) -> Dict[STKey, torch.Tensor]:
-    result = {}
-    for key, value in d.items():
-        if isinstance(value, torch.Tensor):
-            result[STKey((key,), False)] = value
-        elif isinstance(value, dict):
-            value = flatten_dict(value)
-            for inner_key, inner_value in value.items():
-                result[STKey((key,) + inner_key.keys, inner_key.value_is_pickled)] = inner_value
-        else:
-            pickled = bytearray(pickle.dumps(value))
-            pickled_tensor = torch.frombuffer(pickled, dtype=torch.uint8)
-            result[STKey((key,), True)] = pickled_tensor
-    return result
-
-
-def unflatten_dict(d: Dict[STKey, torch.Tensor]) -> Dict:
-    result: Dict = {}
-
-    for key, value in d.items():
-        if key.value_is_pickled:
-            value = pickle.loads(value.numpy().data)
-
-        target_dict = result
-        for k in key.keys[:-1]:
-            new_target_dict = target_dict.get(k)
-            if new_target_dict is None:
-                new_target_dict = {}
-                target_dict[k] = new_target_dict
-            target_dict = new_target_dict
-        target_dict[key.keys[-1]] = value
-
-    return result
-
-
-def state_dict_to_safetensors_file(state_dict: Dict, filename: PathOrStr):
-    state_dict = flatten_dict(state_dict)
-    state_dict = {encode_key(k): v for k, v in state_dict.items()}
-    safetensors.torch.save_file(state_dict, filename)
-
-
-def safetensors_file_to_state_dict(filename: PathOrStr, map_location: Optional[str] = None) -> Dict:
-    if map_location is None:
-        map_location = "cpu"
-    state_dict = safetensors.torch.load_file(filename, device=map_location)
-    state_dict = {decode_key(k): v for k, v in state_dict.items()}
-    return unflatten_dict(state_dict)