gguf: Add ability to load local file

packages/gguf/package.json

@@ -49,6 +49,7 @@
 	"author": "Hugging Face",
 	"license": "MIT",
 	"devDependencies": {
+		"@types/node": "^20.12.8",
 		"type-fest": "^3.9.0"
 	}
 }

packages/gguf/src/gguf.spec.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it } from "vitest";
 import { GGMLQuantizationType, gguf, ggufAllShards, parseGgufShardFilename } from "./gguf";
+import fs from "node:fs";
 
 const URL_LLAMA = "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/191239b/llama-2-7b-chat.Q2_K.gguf";
 const URL_MISTRAL_7B =
@@ -223,6 +224,19 @@ describe("gguf", () => {
 		});
 	});
 
+	it("should parse a local file", async () => {
+		// download the file and save to .cache folder
+		if (!fs.existsSync(".cache")) {
+			fs.mkdirSync(".cache");
+		}
+		const res = await fetch(URL_V1);
+		const arrayBuf = await res.arrayBuffer();
+		fs.writeFileSync(".cache/model.gguf", Buffer.from(arrayBuf));
+
+		const { metadata } = await gguf(".cache/model.gguf", { localFile: true });
+		expect(metadata["general.name"]).toEqual("tinyllamas-stories-260k");
+	});
+
 	it("should detect sharded gguf filename", async () => {
 		const ggufPath = "grok-1/grok-1-q4_0-00003-of-00009.gguf"; // https://huggingface.co/ggml-org/models/blob/fcf344adb9686474c70e74dd5e55465e9e6176ef/grok-1/grok-1-q4_0-00003-of-00009.gguf
 		const ggufShardFileInfo = parseGgufShardFilename(ggufPath);

packages/gguf/src/gguf.ts

@@ -49,7 +49,7 @@ const HTTP_TOTAL_MAX_SIZE = 50 * 10 ** 6; /// 50MB
  * Internal stateful instance to fetch ranges of HTTP data when needed
  */
 class RangeView {
-	private chunk: number;
+	protected chunk: number;
 	private buffer: ArrayBuffer;
 	private dataView: DataView;
 
@@ -58,7 +58,7 @@ class RangeView {
 	}
 
 	constructor(
-		public url: string,
+		public uri: string,
 		private params?: {
 			/**
 			 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
@@ -81,7 +81,7 @@ class RangeView {
 		const range = [this.chunk * HTTP_CHUNK_SIZE, (this.chunk + 1) * HTTP_CHUNK_SIZE - 1];
 		const buf = new Uint8Array(
 			await (
-				await (this.params?.fetch ?? fetch)(this.url, {
+				await (this.params?.fetch ?? fetch)(this.uri, {
 					headers: {
 						...(this.params?.additionalFetchHeaders ?? {}),
 						Range: `bytes=${range[0]}-${range[1]}`,
@@ -128,6 +128,38 @@ class RangeView {
 	}
 }
 
+/**
+ * Internal stateful instance to read ranges of local file when needed.
+ * Only usable in with nodejs FS API.
+ */
+class RangeViewLocalFile extends RangeView {
+	/**
+	 * Read a new chunk from local file system.
+	 */
+	override fetchChunk(): Promise<void> {
+		return new Promise<void>((resolve, reject) => {
+			const Buffer = global.Buffer;
+			if (typeof Buffer === "undefined") {
+				reject(new Error("localFile cannot be used in browser"));
+				return;
+			}
+			let buffer = Buffer.alloc(0);
+			// TODO: using global "import" will make the whole script fails to load on browser
+			// eslint-disable-next-line @typescript-eslint/no-var-requires
+			const stream = require("node:fs").createReadStream(this.uri, {
+				start: this.chunk * HTTP_CHUNK_SIZE,
+				end: (this.chunk + 1) * HTTP_CHUNK_SIZE,
+			});
+			stream.on("error", reject);
+			stream.on("data", (chunk: Buffer) => (buffer = Buffer.concat([buffer, chunk])));
+			stream.on("end", () => {
+				this.appendBuffer(buffer);
+				resolve();
+			});
+		});
+	}
+}
+
 interface Slice<T> {
 	value: T;
 	length: number;
@@ -205,38 +237,41 @@ function readMetadataValue(
 }
 
 export async function gguf(
-	url: string,
+	uri: string,
 	params: {
 		/**
 		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 		 */
 		fetch?: typeof fetch;
 		additionalFetchHeaders?: Record<string, string>;
 		computeParametersCount: true;
+		localFile?: boolean;
 	}
 ): Promise<GGUFParseOutput & { parameterCount: number }>;
 export async function gguf(
-	url: string,
+	uri: string,
 	params?: {
 		/**
 		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 		 */
 		fetch?: typeof fetch;
 		additionalFetchHeaders?: Record<string, string>;
+		localFile?: boolean;
 	}
 ): Promise<GGUFParseOutput>;
 export async function gguf(
-	url: string,
+	uri: string,
 	params?: {
 		/**
 		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 		 */
 		fetch?: typeof fetch;
 		additionalFetchHeaders?: Record<string, string>;
 		computeParametersCount?: boolean;
+		localFile?: boolean;
 	}
 ): Promise<GGUFParseOutput & { parameterCount?: number }> {
-	const r = new RangeView(url, params);
+	const r = params?.localFile ? new RangeViewLocalFile(uri, params) : new RangeView(uri, params);
 	await r.fetchChunk();
 
 	const checkBuffer = (buffer: Uint8Array, header: Uint8Array) => {
@@ -376,7 +411,7 @@ export async function ggufAllShards(
 
 		const PARALLEL_DOWNLOADS = 20;
 		const shards = await promisesQueue(
-			urls.map((shardUrl) => () => gguf(shardUrl, { computeParametersCount: true })),
+			urls.map((shardUrl) => () => gguf(shardUrl, { ...params, computeParametersCount: true })),
 			PARALLEL_DOWNLOADS
 		);
 		return {

packages/gguf/tsup.config.ts

@@ -21,6 +21,8 @@ const browserConfig: Options = {
 	target: "es2018",
 	splitting: true,
 	outDir: "dist/browser",
+	// We specify external libs only to be able to build. We're not using them on browser.
+	external: ["node:fs"],
 };
 
 export default [nodeConfig, browserConfig];