shortcut like crazy

Author: julien-c

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

@@ -141,7 +141,7 @@ async function parseShardedIndex(
 		 */
 		fetch?: typeof fetch;
 	} & Partial<CredentialsParams>
-): Promise<{ index: SafetensorsIndexJson; headers: SafetensorsShardedHeaders }> {
+): Promise<SafetensorsIndexJson> {
 	const indexBlob = await downloadFile({
 		...params,
 		path,
@@ -151,14 +151,28 @@ async function parseShardedIndex(
 		throw new SafetensorParseError(`Failed to parse file ${path}: failed to fetch safetensors index.`);
 	}
 
-	// no validation for now, we assume it's a valid IndexJson.
-	let index: SafetensorsIndexJson;
 	try {
-		index = JSON.parse(await indexBlob.slice(0, 10_000_000).text());
+		// no validation for now, we assume it's a valid IndexJson.
+		const index = JSON.parse(await indexBlob.slice(0, 10_000_000).text());
+		return index;
 	} catch (error) {
 		throw new SafetensorParseError(`Failed to parse file ${path}: not a valid JSON.`);
 	}
+}
 
+async function fetchAllHeaders(
+	path: string,
+	index: SafetensorsIndexJson,
+	params: {
+		repo: RepoDesignation;
+		revision?: string;
+		hubUrl?: string;
+		/**
+		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
+		 */
+		fetch?: typeof fetch;
+	} & Partial<CredentialsParams>
+): Promise<SafetensorsShardedHeaders> {
 	const pathPrefix = path.slice(0, path.lastIndexOf("/") + 1);
 	const filenames = [...new Set(Object.values(index.weight_map))];
 	const shardedMap: SafetensorsShardedHeaders = Object.fromEntries(
@@ -170,7 +184,7 @@ async function parseShardedIndex(
 			PARALLEL_DOWNLOADS
 		)
 	);
-	return { index, headers: shardedMap };
+	return shardedMap;
 }
 
 /**
@@ -191,6 +205,7 @@ export async function parseSafetensorsMetadata(
 		 * @default false
 		 */
 		computeParametersCount: true;
+		fetchAllHeaders?: boolean;
 		hubUrl?: string;
 		revision?: string;
 		/**
@@ -210,6 +225,12 @@ export async function parseSafetensorsMetadata(
 		 * @default false
 		 */
 		computeParametersCount?: boolean;
+		/**
+		 * Always fetch all headers (no shortcut)
+		 *
+		 * @default false
+		 */
+		fetchAllHeaders?: boolean;
 		hubUrl?: string;
 		revision?: string;
 		/**
@@ -223,6 +244,7 @@ export async function parseSafetensorsMetadata(
 		repo: RepoDesignation;
 		path?: string;
 		computeParametersCount?: boolean;
+		fetchAllHeaders?: boolean;
 		hubUrl?: string;
 		revision?: string;
 		/**
@@ -255,15 +277,31 @@ export async function parseSafetensorsMetadata(
 		(params.path && RE_SAFETENSORS_INDEX_FILE.test(params.path)) ||
 		(await fileExists({ ...params, path: SAFETENSORS_INDEX_FILE }))
 	) {
-		const { index, headers } = await parseShardedIndex(params.path ?? SAFETENSORS_INDEX_FILE, params);
+		const path = params.path ?? SAFETENSORS_INDEX_FILE;
+		const index = await parseShardedIndex(path, params);
+
+		const shardedMap =
+			params.fetchAllHeaders || (params.computeParametersCount && !index.metadata?.total_parameters)
+				? await fetchAllHeaders(path, index, params)
+				: {};
+
+		if (params.computeParametersCount && index.metadata?.total_parameters) {
+			/// shortcut: get param count directly from metadata
+			return {
+				sharded: true,
+				index,
+				headers: shardedMap,
+				parameterCount: { UNK: parseInt(index.metadata.total_parameters.toString()) },
+			};
+		}
 
 		return {
 			sharded: true,
 			index,
-			headers,
+			headers: shardedMap,
 			...(params.computeParametersCount
 				? {
-						parameterCount: computeNumOfParamsByDtypeSharded(index, headers),
+						parameterCount: computeNumOfParamsByDtypeSharded(shardedMap),
 				  }
 				: undefined),
 		};
@@ -289,14 +327,7 @@ function computeNumOfParamsByDtypeSingleFile(header: SafetensorsFileHeader): Par
 	return counter;
 }
 
-function computeNumOfParamsByDtypeSharded(
-	index: SafetensorsIndexJson,
-	shardedMap: SafetensorsShardedHeaders
-): Partial<Record<Dtype, number>> {
-	if (index.metadata?.total_parameters) {
-		/// shortcut: get param count directly from metadata
-		return { UNK: parseInt(index.metadata.total_parameters.toString()) };
-	}
+function computeNumOfParamsByDtypeSharded(shardedMap: SafetensorsShardedHeaders): Partial<Record<Dtype, number>> {
 	const counter: Partial<Record<Dtype, number>> = {};
 	for (const header of Object.values(shardedMap)) {
 		for (const [k, v] of typedEntries(computeNumOfParamsByDtypeSingleFile(header))) {