extract_keys.py

import torch
import safetensors.torch
import sys

def dump_keys(parent, suffix=''):
    for k in sorted(parent.keys()):
        if isinstance(parent[k], torch.Tensor):
            print(f'{suffix}{k} {list(parent[k].shape)}')
        else:
            dump_keys(parent[k], f'{suffix}{k}.')

if __name__ == '__main__':
    fn = sys.argv[1]
    if fn.endswith('.safetensors'):
        checkpoint = safetensors.torch.load_file(fn)
    else:
        checkpoint = torch.load(fn)
    dump_keys(checkpoint)