Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| import os | |
| import sys | |
| import json | |
| from rich import print # pylint: disable=redefined-builtin | |
| def read_metadata(fn): | |
| res = {} | |
| with open(fn, mode="rb") as f: | |
| metadata_len = f.read(8) | |
| metadata_len = int.from_bytes(metadata_len, "little") | |
| json_start = f.read(2) | |
| if metadata_len <= 2 or json_start not in (b'{"', b"{'"): | |
| print(f"Not a valid safetensors file: {fn}") | |
| json_data = json_start + f.read(metadata_len-2) | |
| json_obj = json.loads(json_data) | |
| for k, v in json_obj.get("__metadata__", {}).items(): | |
| res[k] = v | |
| if isinstance(v, str) and v[0:1] == '{': | |
| try: | |
| res[k] = json.loads(v) | |
| except Exception: | |
| pass | |
| print(f"{fn}: {json.dumps(res, indent=4)}") | |
| def main(): | |
| if len(sys.argv) == 0: | |
| print('metadata:', 'no files specified') | |
| for fn in sys.argv: | |
| if os.path.isfile(fn): | |
| read_metadata(fn) | |
| elif os.path.isdir(fn): | |
| for root, _dirs, files in os.walk(fn): | |
| for file in files: | |
| read_metadata(os.path.join(root, file)) | |
| if __name__ == '__main__': | |
| sys.argv.pop(0) | |
| main() | |