Return file paths from gguf and safetensors parsers

packages/gguf/src/gguf.spec.ts

@@ -319,8 +319,17 @@ describe("gguf", () => {
 	});
 
 	it("should get param count for sharded gguf", async () => {
-		const { parameterCount } = await ggufAllShards(URL_SHARDED_GROK);
+		const { parameterCount, urls } = await ggufAllShards(URL_SHARDED_GROK);
 		expect(parameterCount).toEqual(316_490_127_360); // 316B
+		expect(urls).toHaveLength(9);
+		expect(urls[0]).toMatch(/00001-of-00009\.gguf$/);
+		expect(urls[8]).toMatch(/00009-of-00009\.gguf$/);
+	});
+
+	it("should return urls for single (non-sharded) gguf", async () => {
+		const { urls } = await ggufAllShards(URL_LLAMA);
+		expect(urls).toHaveLength(1);
+		expect(urls[0]).toEqual(URL_LLAMA);
 	});
 
 	it("parse quant label", async () => {

packages/gguf/src/gguf.ts

@@ -818,7 +818,11 @@ export async function ggufAllShards(
 		parallelDownloads?: number;
 		allowLocalFile?: boolean;
 	},
-): Promise<{ shards: GGUFParseOutput[]; parameterCount: number }> {
+): Promise<{
+	shards: GGUFParseOutput[];
+	parameterCount: number;
+	urls: string[];
+}> {
 	const parallelDownloads = params?.parallelDownloads ?? PARALLEL_DOWNLOADS;
 	if (parallelDownloads < 1) {
 		throw new TypeError("parallelDownloads must be greater than 0");
@@ -840,6 +844,7 @@ export async function ggufAllShards(
 		return {
 			shards,
 			parameterCount: shards.map(({ parameterCount }) => parameterCount).reduce((acc, val) => acc + val, 0),
+			urls,
 		};
 	} else {
 		const { metadata, tensorInfos, tensorDataOffset, littleEndian, parameterCount, tensorInfoByteRange } = await gguf(
@@ -849,6 +854,10 @@ export async function ggufAllShards(
 				computeParametersCount: true,
 			},
 		);
-		return { shards: [{ metadata, tensorInfos, tensorDataOffset, littleEndian, tensorInfoByteRange }], parameterCount };
+		return {
+			shards: [{ metadata, tensorInfos, tensorDataOffset, littleEndian, tensorInfoByteRange }],
+			parameterCount,
+			urls: [url],
+		};
 	}
 }

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

@@ -29,6 +29,8 @@ describe("parseSafetensorsMetadata", () => {
 		assert.deepStrictEqual(parse.parameterCount, { F32: 110_106_428 });
 		assert.deepStrictEqual(sum(Object.values(parse.parameterCount)), 110_106_428);
 		// total params = 110m
+
+		assert.deepStrictEqual(parse.filepaths, ["model.safetensors"]);
 	});
 
 	it("fetch info for sharded (with the default conventional filename)", async () => {
@@ -54,6 +56,10 @@ describe("parseSafetensorsMetadata", () => {
 		assert.deepStrictEqual(parse.parameterCount, { BF16: 176_247_271_424 });
 		assert.deepStrictEqual(sum(Object.values(parse.parameterCount)), 176_247_271_424);
 		// total params = 176B
+
+		assert.strictEqual(parse.filepaths[0], "model.safetensors.index.json");
+		assert.strictEqual(parse.filepaths.length, 73); // 1 index + 72 shards
+		assert.ok(parse.filepaths.includes("model_00012-of-00072.safetensors"));
 	});
 
 	it("fetch info for single-file with multiple dtypes", async () => {
@@ -91,6 +97,8 @@ describe("parseSafetensorsMetadata", () => {
 
 		assert.deepStrictEqual(parse.parameterCount, { F32: 859_520_964 });
 		assert.deepStrictEqual(sum(Object.values(parse.parameterCount)), 859_520_964);
+
+		assert.deepStrictEqual(parse.filepaths, ["unet/diffusion_pytorch_model.safetensors"]);
 	});
 
 	it("fetch info for sharded with file path", async () => {
@@ -113,6 +121,11 @@ describe("parseSafetensorsMetadata", () => {
 
 		assert.deepStrictEqual(parse.parameterCount, { BF16: 8_537_680_896 });
 		assert.deepStrictEqual(sum(Object.values(parse.parameterCount)), 8_537_680_896);
+
+		assert.strictEqual(parse.filepaths[0], "7b/1/model.safetensors.index.json");
+		assert.strictEqual(parse.filepaths.length, 5); // 1 index + 4 shards
+		assert.ok(parse.filepaths.includes("7b/1/model-00001-of-00004.safetensors"));
+		assert.ok(parse.filepaths.includes("7b/1/model-00004-of-00004.safetensors"));
 	});
 
 	it("fetch info for sharded, but get param count directly from metadata", async () => {

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

@@ -93,13 +93,15 @@ export type SafetensorsParseFromRepo =
 			header: SafetensorsFileHeader;
 			parameterCount?: Partial<Record<Dtype, number>>;
 			parameterTotal?: number;
+			filepaths: string[];
 	  }
 	| {
 			sharded: true;
 			index: SafetensorsIndexJson;
 			headers: SafetensorsShardedHeaders;
 			parameterCount?: Partial<Record<Dtype, number>>;
 			parameterTotal?: number;
+			filepaths: string[];
 	  };
 
 /**
@@ -323,14 +325,20 @@ export async function parseSafetensorsMetadata(
 					parameterTotal: parseTotalParameters(header.__metadata__.total_parameters),
 				}
 			: undefined;
-		return { sharded: false, header, ...paramStats };
+		return {
+			sharded: false,
+			header,
+			...paramStats,
+			filepaths: [params.path ?? SAFETENSORS_FILE],
+		};
 	} else if (
 		(params.path && RE_SAFETENSORS_INDEX_FILE.test(params.path)) ||
 		(await fileExists({ ...params, path: SAFETENSORS_INDEX_FILE }))
 	) {
 		const path = params.path ?? SAFETENSORS_INDEX_FILE;
 		const index = await parseShardedIndex(path, params);
 		const shardedMap = await fetchAllHeaders(path, index, params);
+		const pathPrefix = path.slice(0, path.lastIndexOf("/") + 1);
 
 		const paramStats = params.computeParametersCount
 			? {
@@ -339,7 +347,13 @@ export async function parseSafetensorsMetadata(
 					parameterTotal: parseTotalParameters(index.metadata?.total_parameters),
 				}
 			: undefined;
-		return { sharded: true, index, headers: shardedMap, ...paramStats };
+		return {
+			sharded: true,
+			index,
+			headers: shardedMap,
+			...paramStats,
+			filepaths: [path, ...Object.keys(shardedMap).map((filename) => pathPrefix + filename)],
+		};
 	} else {
 		throw new Error("model id does not seem to contain safetensors weights");
 	}