| 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 个键名预览:")
|
|
|
|
|
| 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} 个键名")
|
|
|
|
|
| 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()
|
|
|
|
|
| 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) |
