Skip to content

Commit a22b767

Browse files
authored
Allow to customize parallelDownload for ggufAllShards (#859)
Could be related to https://huggingface.slack.com/archives/C02EMARJ65P/p1724166668261449
1 parent 207340f commit a22b767

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

packages/gguf/src/gguf.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ export { GGUF_QUANT_DESCRIPTIONS } from "./quant-descriptions";
99

1010
export const RE_GGUF_FILE = /\.gguf$/;
1111
export const RE_GGUF_SHARD_FILE = /^(?<prefix>.*?)-(?<shard>\d{5})-of-(?<total>\d{5})\.gguf$/;
12+
const PARALLEL_DOWNLOADS = 20;
1213

1314
export interface GgufShardFileInfo {
1415
prefix: string;
@@ -401,8 +402,13 @@ export async function ggufAllShards(
401402
*/
402403
fetch?: typeof fetch;
403404
additionalFetchHeaders?: Record<string, string>;
405+
parallelDownloads?: number;
404406
}
405407
): Promise<{ shards: GGUFParseOutput[]; parameterCount: number }> {
408+
const parallelDownloads = params?.parallelDownloads ?? PARALLEL_DOWNLOADS;
409+
if (parallelDownloads < 1) {
410+
throw new TypeError("parallelDownloads must be greater than 0");
411+
}
406412
const ggufShardFileInfo = parseGgufShardFilename(url);
407413
if (ggufShardFileInfo) {
408414
const total = parseInt(ggufShardFileInfo.total);
@@ -413,10 +419,9 @@ export async function ggufAllShards(
413419
urls.push(`${prefix}-${shardIdx.toString().padStart(5, "0")}-of-${total.toString().padStart(5, "0")}.gguf`);
414420
}
415421

416-
const PARALLEL_DOWNLOADS = 20;
417422
const shards = await promisesQueue(
418423
urls.map((shardUrl) => () => gguf(shardUrl, { ...params, computeParametersCount: true })),
419-
PARALLEL_DOWNLOADS
424+
parallelDownloads
420425
);
421426
return {
422427
shards,

0 commit comments

Comments
 (0)