From b8279c82d0783eb407250414078ed2df2acbb9bb Mon Sep 17 00:00:00 2001 From: ubik2 Date: Tue, 2 May 2023 02:47:35 -0700 Subject: [PATCH] When loading a safetensors file, ignore the metadata header; If no pt or pth files are available, attempt to load safetensors files --- convert.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/convert.py b/convert.py index 7f7ae05fa..75eeb6125 100644 --- a/convert.py +++ b/convert.py @@ -750,7 +750,7 @@ def lazy_load_safetensors_file(fp: IO[bytes], path: Path) -> ModelPlus: return UnquantizedTensor(np.frombuffer(buf, dtype=numpy_dtype).reshape(shape)) description = f'safetensors begin={begin} end={end} type={data_type} path={path}' return LazyTensor(load, shape, data_type, description) - model = {name: convert(info) for (name, info) in header.items()} + model = {name: convert(info) for (name, info) in header.items() if name != '__metadata__'} return ModelPlus(model=model, paths=[path], format='safetensors', vocab=None) @@ -1037,6 +1037,10 @@ def load_some_model(path: Path) -> ModelPlus: if path.is_dir(): globs = ["consolidated.00.pth", "pytorch_model-00001-of-*.bin", "*.pt"] files = [file for glob in globs for file in path.glob(glob)] + if not files: + # Check if it's a set of safetensors files + globs = ["model-00001-of-*.safetensors"] + files = [file for glob in globs for file in path.glob(glob)] if not files: # Try GGML too, but with lower priority, since if both a non-GGML # model and a GGML model exist in the same directory, we assume the