[safetensors] add named groups to RE_SAFETENSORS_SHARD_FILE regex (#622)

add named groups to RE_SAFETENSORS_SHARD_FILE regex

safetensors equivalent of
#621

---------

Co-authored-by: Eliott C. <[email protected]>

Author: Mishig

packages/hub/src/lib/parse-safetensors-metadata.spec.ts

@@ -1,5 +1,5 @@
 import { assert, it, describe } from "vitest";
-import { parseSafetensorsMetadata } from "./parse-safetensors-metadata";
+import { RE_SAFETENSORS_SHARD_FILE, parseSafetensorsMetadata } from "./parse-safetensors-metadata";
 import { sum } from "../utils/sum";
 
 describe("parseSafetensorsMetadata", () => {
@@ -109,4 +109,14 @@ describe("parseSafetensorsMetadata", () => {
 		assert.deepStrictEqual(parse.parameterCount, { BF16: 8_537_680_896 });
 		assert.deepStrictEqual(sum(Object.values(parse.parameterCount)), 8_537_680_896);
 	});
+
+	it("should detect sharded safetensors filename", async () => {
+		const safetensorsFilename = "model_00005-of-00072.safetensors"; // https://huggingface.co/bigscience/bloom/blob/4d8e28c67403974b0f17a4ac5992e4ba0b0dbb6f/model_00005-of-00072.safetensors
+		const match = safetensorsFilename.match(RE_SAFETENSORS_SHARD_FILE);
+
+		assert.strictEqual(RE_SAFETENSORS_SHARD_FILE.test(safetensorsFilename), true);
+		assert.strictEqual(match?.groups?.prefix, "model");
+		assert.strictEqual(match?.groups?.shard, "00005");
+		assert.strictEqual(match?.groups?.total, "00072");
+	});
 });

packages/hub/src/lib/parse-safetensors-metadata.ts

@@ -14,7 +14,7 @@ export const SAFETENSORS_INDEX_FILE = "model.safetensors.index.json";
 /// but in some situations safetensors weights have different filenames.
 export const RE_SAFETENSORS_FILE = /\.safetensors$/;
 export const RE_SAFETENSORS_INDEX_FILE = /\.safetensors\.index\.json$/;
-export const RE_SAFETENSORS_SHARD_FILE = /\d{5}-of-\d{5}\.safetensors$/;
+export const RE_SAFETENSORS_SHARD_FILE = /^(?<prefix>.*?)[_-]?(?<shard>\d{5})-of-(?<total>\d{5})\.safetensors$/;
 const PARALLEL_DOWNLOADS = 20;
 const MAX_HEADER_LENGTH = 25_000_000;