File size: 2,767 Bytes
9f09426 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import os
import argparse
from safetensors import safe_open
def inspect_safetensors(file_path):
if not os.path.exists(file_path):
print(f"❌ 错误: 找不到文件 '{file_path}'")
return
print(f"🔍 正在读取文件: {file_path} ...")
# 获取输出日志文件名
base_name = os.path.splitext(os.path.basename(file_path))[0]
output_txt = f"{base_name}_keys_list.txt"
try:
with safe_open(file_path, framework="pt", device="cpu") as f:
keys = sorted(list(f.keys()))
total_keys = len(keys)
print(f"📊 文件读取成功！共有 {total_keys} 个权重键名 (Keys)。")
print("-" * 50)
print("💡 前 20 个键名预览:")
# 打印前20个预览
for i, key in enumerate(keys[:20]):
tensor_slice = f.get_tensor(key)
print(f" [{i+1}] {key} --> shape: {list(tensor_slice.shape)} | dtype: {tensor_slice.dtype}")
if total_keys > 20:
print(f" ... 以及其余 {total_keys - 20} 个键名")
# 将所有键名导出到 txt 文件
with open(output_txt, "w", encoding="utf-8") as txt_f:
txt_f.write(f"模型文件: {file_path}\n")
txt_f.write(f"总计键名数: {total_keys}\n")
txt_f.write("=" * 60 + "\n")
for key in keys:
tensor_slice = f.get_tensor(key)
txt_f.write(f"{key}\t{list(tensor_slice.shape)}\t{tensor_slice.dtype}\n")
print("-" * 50)
print(f"💾 完整键名列表已导出至当前目录: \n👉 {output_txt}")
except Exception as e:
print(f"❌ 读取过程中发生错误: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="快速检查 SafeTensors 模型内部的全部 Key 名")
parser.add_argument("file_path", type=str, nargs="?", help="输入的 .safetensors 文件路径")
args = parser.parse_args()
# 如果运行脚本时没有传参，则自动在当前目录下找第一个 safetensors 文件
if not args.file_path:
current_dir_files = [f for f in os.listdir(".") if f.endswith(".safetensors")]
if current_dir_files:
print(f"提示: 未指定输入文件，自动选择当前目录下找到的: {current_dir_files[0]}")
inspect_safetensors(current_dir_files[0])
else:
print("❌ 错误: 当前目录下没有找到任何 .safetensors 文件，请在后面加上文件路径参数。")
else:
inspect_safetensors(args.file_path) |