File size: 1,423 Bytes
64ddf8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import json
import os
from safetensors import safe_open
def get_tensor(matrix_name, base_dir, index_json="model.safetensors.index.json"):
fp = os.path.join(base_dir, index_json)
with open(fp, "r") as fh:
index = json.load(fh)
try:
st_file = index["weight_map"][matrix_name]
mp = os.path.join(base_dir, st_file)
with safe_open(mp, framework="pt", device="cpu") as f:
print(f.keys())
return f.get_tensor(matrix_name)
except Exception:
raise ValueError(f"Invalid key {matrix_name}")
def get_tensor2(matrix_name, st_file_fp):
try:
with safe_open(st_file_fp, framework="pt", device="cpu") as f:
return f.get_tensor(matrix_name)
except Exception:
raise ValueError("Fail to retrieve tensor keys")
def get_tensor_keys(base_dir, st_file):
try:
mp = os.path.join(base_dir, st_file)
with safe_open(mp, framework="pt", device="cpu") as f:
return f.keys()
except Exception:
raise ValueError("Fail to retrieve tensor keys")
def get_tensor_dual(matrix1, matrix2, base_dir, st_file):
try:
mp = os.path.join(base_dir, st_file)
with safe_open(mp, framework="pt", device="cpu") as f:
return f.get_tensor(matrix1), f.get_tensor(matrix2)
except Exception:
raise ValueError(f"Invalid key {matrix1}")
|