From ce53e31f0ac3eea4a6a70443ddab4f1e759c7b60 Mon Sep 17 00:00:00 2001 From: Duane Moody Date: Fri, 1 Dec 2023 14:06:22 -0700 Subject: [PATCH] Add files via upload Proof-of-concept utility for retrieving a named JSON object from `__metadata__` by passing it and the safetensors filename. --- safetensors_m.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 safetensors_m.py diff --git a/safetensors_m.py b/safetensors_m.py new file mode 100644 index 0000000..ba6d691 --- /dev/null +++ b/safetensors_m.py @@ -0,0 +1,22 @@ +import json, sys +from safetensors_file import SafeTensorsFile +from safetensors_worker import _ParseMore + +def get_object(tensorsfile: str) -> str: + with SafeTensorsFile.open_file(tensorsfile, quiet=True) as s: + js = s.get_header() + + if "__metadata__" not in js: + return "Error: File header does not contain a `__metadata__` item" + md = js["__metadata__"] + if md_object not in md: + return f'Error: Metadata does not contain a `{md_object}` item, did you spell it right?' + _ParseMore(md) # pretty print the metadata + stf = md[md_object] + return json.dumps(stf, ensure_ascii=False, separators=(', ', ': '), indent=4) + +md_object = sys.argv[1] +tensorsfile = sys.argv[2] +hdata = get_object(tensorsfile) + +print(hdata)