Skip to content

Commit

Permalink
Allow to customize parallelDownload for ggufAllShards (#859)
Browse files Browse the repository at this point in the history
  • Loading branch information
Kakulukian authored Aug 26, 2024
1 parent 207340f commit a22b767
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions packages/gguf/src/gguf.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ export { GGUF_QUANT_DESCRIPTIONS } from "./quant-descriptions";

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

export interface GgufShardFileInfo {
prefix: string;
Expand Down Expand Up @@ -401,8 +402,13 @@ export async function ggufAllShards(
*/
fetch?: typeof fetch;
additionalFetchHeaders?: Record<string, string>;
parallelDownloads?: number;
}
): Promise<{ shards: GGUFParseOutput[]; parameterCount: number }> {
const parallelDownloads = params?.parallelDownloads ?? PARALLEL_DOWNLOADS;
if (parallelDownloads < 1) {
throw new TypeError("parallelDownloads must be greater than 0");
}
const ggufShardFileInfo = parseGgufShardFilename(url);
if (ggufShardFileInfo) {
const total = parseInt(ggufShardFileInfo.total);
Expand All @@ -413,10 +419,9 @@ export async function ggufAllShards(
urls.push(`${prefix}-${shardIdx.toString().padStart(5, "0")}-of-${total.toString().padStart(5, "0")}.gguf`);
}

const PARALLEL_DOWNLOADS = 20;
const shards = await promisesQueue(
urls.map((shardUrl) => () => gguf(shardUrl, { ...params, computeParametersCount: true })),
PARALLEL_DOWNLOADS
parallelDownloads
);
return {
shards,
Expand Down

0 comments on commit a22b767

Please sign in to comment.