Commit 8dc92022 authored by Brendan Hoar's avatar Brendan Hoar Committed by GitHub

Better error handling when unable to read metadata from safetensors file

parent 3902aa22
...@@ -280,10 +280,12 @@ def read_metadata_from_safetensors(filename): ...@@ -280,10 +280,12 @@ def read_metadata_from_safetensors(filename):
json_start = file.read(2) json_start = file.read(2)
assert metadata_len > 2 and json_start in (b'{"', b"{'"), f"{filename} is not a safetensors file" assert metadata_len > 2 and json_start in (b'{"', b"{'"), f"{filename} is not a safetensors file"
json_data = json_start + file.read(metadata_len-2)
json_obj = json.loads(json_data)
res = {} res = {}
try:
json_data = json_start + file.read(metadata_len-2)
json_obj = json.loads(json_data)
for k, v in json_obj.get("__metadata__", {}).items(): for k, v in json_obj.get("__metadata__", {}).items():
res[k] = v res[k] = v
if isinstance(v, str) and v[0:1] == '{': if isinstance(v, str) and v[0:1] == '{':
...@@ -291,6 +293,8 @@ def read_metadata_from_safetensors(filename): ...@@ -291,6 +293,8 @@ def read_metadata_from_safetensors(filename):
res[k] = json.loads(v) res[k] = json.loads(v)
except Exception: except Exception:
pass pass
except:
errors.report(f"Error reading metadata from file: {filename}", exc_info=True)
return res return res
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment