XetBlob util (#1257)

Python lib code:
https://github.com/huggingface-internal/xetpoc_huggingface_hub/pull/14/files
(cc @bpronan @Wauplin @hanouticelina)

# Todo

Add more test cases

- [x] on a larger amount of data (need to fetch several terms)
- [x] when offset_into_first_range > a chunk length
- [ ] when term.range.start > fetchinfo.range.start => created
#1285
- [x] check range request when requesting range at exact chunk
boundaries, or just before/just after

Author: coyotte508

.github/workflows/test.yml

@@ -61,11 +61,13 @@ jobs:
             echo "SINCE=$(git merge-base origin/${{ github.event.pull_request.base.ref }} ${{ github.sha }})" >> $GITHUB_OUTPUT
           fi
 
+      - run: google-chrome --version
+
       - run: npm install -g corepack@latest && corepack enable
 
       - uses: actions/setup-node@v3
         with:
-          node-version: "20"
+          node-version: "22"
           cache: "pnpm"
           cache-dependency-path: "**/pnpm-lock.yaml"
       - run: |

CONTRIBUTING.md

@@ -18,7 +18,11 @@ It's not a hard requirement, but please consider using an icon from [Gitmoji](ht
 
 ## Tests
 
-If you want to run only specific tests, you can do `pnpm test -- -t "test name"`
+If you want to run only specific tests, you can do `pnpm test -- -t "test name"`.
+
+You can also do `npx vitest ./packages/hub/src/utils/XetBlob.spec.ts` to run a specific test file.
+
+Or `cd packages/hub && npx vitest --browser.name=chrome --browser.headless --config vitest-browser.config.mts ./src/utils/XetBlob.spec.ts` to run browser tests on a specific file
 
 ## Adding a package
 

packages/hub/.eslintignore

@@ -1,2 +1,3 @@
 dist
 sha256.js
+src/vendor

packages/hub/README.md

@@ -174,3 +174,4 @@ Under the hood, `@huggingface/hub` uses a lazy blob implementation to load the f
 ## Dependencies
 
 - `@huggingface/tasks` : Typings only
+- `@huggingface/lz4` : URL join utility

packages/hub/src/index.ts

@@ -22,3 +22,4 @@ export { HubApiError, InvalidApiResponseFormatError } from "./error";
  * Only exported for E2Es convenience
  */
 export { sha256 as __internal_sha256 } from "./utils/sha256";
+export { XetBlob as __internal_XetBlob } from "./utils/XetBlob";

packages/hub/src/lib/list-files.spec.ts

@@ -60,7 +60,7 @@ describe("listFiles", () => {
 				type: "file",
 			},
 		]);
-	});
+	}, 30_000);
 
 	it("should fetch the list of files from the repo, including last commit", async () => {
 		const cursor = listFiles({
@@ -146,7 +146,7 @@ describe("listFiles", () => {
 				type: "file",
 			},
 		]);
-	});
+	}, 30_000);
 
 	it("should fetch the list of files from the repo, including subfolders", async () => {
 		const cursor = listFiles({
@@ -165,5 +165,5 @@ describe("listFiles", () => {
 		}
 
 		assert(files.some((file) => file.path === "data/XSUM-EMNLP18-Summary-Data-Original.tar.gz"));
-	});
+	}, 30_000);
 });

packages/hub/src/utils/WebBlob.spec.ts

@@ -58,7 +58,7 @@ describe("WebBlob", () => {
 		expect(webBlob).toBeInstanceOf(WebBlob);
 		expect(webBlob).toMatchObject({ url });
 		expect(await webBlob.slice(10, 22).text()).toBe("__metadata__");
-	});
+	}, 30_000);
 
 	it("should lazy load a Xet file hosted on Hugging Face", async () => {
 		const stableDiffusionUrl =
@@ -70,7 +70,7 @@ describe("WebBlob", () => {
 		expect(webBlob).toBeInstanceOf(WebBlob);
 		expect(webBlob).toMatchObject({ url });
 		expect(await webBlob.slice(10, 22).text()).toBe("__metadata__");
-	});
+	}, 30_000);
 
 	it("should create a slice on the file", async () => {
 		const expectedText = fullText.slice(10, 20);

packages/hub/src/utils/WebBlob.ts

@@ -60,7 +60,7 @@ export class WebBlob extends Blob {
 
 	override slice(start = 0, end = this.size): WebBlob {
 		if (start < 0 || end < 0) {
-			new TypeError("Unsupported negative start/end on FileBlob.slice");
+			new TypeError("Unsupported negative start/end on WebBlob.slice");
 		}
 
 		const slice = new WebBlob(

packages/hub/src/utils/XetBlob.spec.ts

@@ -0,0 +1,185 @@
+import { describe, expect, it } from "vitest";
+import { bg4_regoup_bytes, XetBlob } from "./XetBlob";
+
+describe("XetBlob", () => {
+	it("should lazy load the first 22 bytes", async () => {
+		const blob = new XetBlob({
+			repo: {
+				type: "model",
+				name: "celinah/xet-experiments",
+			},
+			hash: "7b3b6d07673a88cf467e67c1f7edef1a8c268cbf66e9dd9b0366322d4ab56d9b",
+			size: 5_234_139_343,
+		});
+
+		expect(await blob.slice(10, 22).text()).toBe("__metadata__");
+	}, 30_000);
+
+	it("should load the first chunk correctly", async () => {
+		let xorbCount = 0;
+		const blob = new XetBlob({
+			repo: {
+				type: "model",
+				name: "celinah/xet-experiments",
+			},
+			hash: "7b3b6d07673a88cf467e67c1f7edef1a8c268cbf66e9dd9b0366322d4ab56d9b",
+			size: 5_234_139_343,
+			fetch: async (url, opts) => {
+				if (typeof url === "string" && url.includes("/xorbs/")) {
+					xorbCount++;
+				}
+				return fetch(url, opts);
+			},
+		});
+
+		const xetDownload = await blob.slice(0, 29928).arrayBuffer();
+		const bridgeDownload = await fetch(
+			"https://huggingface.co/celinah/xet-experiments/resolve/main/model5GB.safetensors",
+			{
+				headers: {
+					Range: "bytes=0-29927",
+				},
+			}
+		).then((res) => res.arrayBuffer());
+
+		expect(new Uint8Array(xetDownload)).toEqual(new Uint8Array(bridgeDownload));
+		expect(xorbCount).toBe(1);
+	}, 30_000);
+
+	it("should load just past the first chunk correctly", async () => {
+		let xorbCount = 0;
+		const blob = new XetBlob({
+			repo: {
+				type: "model",
+				name: "celinah/xet-experiments",
+			},
+			hash: "7b3b6d07673a88cf467e67c1f7edef1a8c268cbf66e9dd9b0366322d4ab56d9b",
+			size: 5_234_139_343,
+			fetch: async (url, opts) => {
+				if (typeof url === "string" && url.includes("/xorbs/")) {
+					xorbCount++;
+				}
+				return fetch(url, opts);
+			},
+		});
+
+		const xetDownload = await blob.slice(0, 29929).arrayBuffer();
+		const bridgeDownload = await fetch(
+			"https://huggingface.co/celinah/xet-experiments/resolve/main/model5GB.safetensors",
+			{
+				headers: {
+					Range: "bytes=0-29928",
+				},
+			}
+		).then((res) => res.arrayBuffer());
+
+		expect(new Uint8Array(xetDownload)).toEqual(new Uint8Array(bridgeDownload));
+		expect(xetDownload.byteLength).toBe(29929);
+		expect(xorbCount).toBe(2);
+	});
+
+	// In github actions, this test doesn't work inside the browser, but it works locally
+	// inside both chrome and chromium browsers
+	// TODO: figure out why
+	if (typeof window === "undefined") {
+		it("should load correctly when loading far into a chunk range", async () => {
+			const blob = new XetBlob({
+				repo: {
+					type: "model",
+					name: "celinah/xet-experiments",
+				},
+				hash: "7b3b6d07673a88cf467e67c1f7edef1a8c268cbf66e9dd9b0366322d4ab56d9b",
+				size: 5_234_139_343,
+			});
+
+			const xetDownload = await blob.slice(10_000_000, 10_100_000).arrayBuffer();
+			const bridgeDownload = await fetch(
+				"https://huggingface.co/celinah/xet-experiments/resolve/main/model5GB.safetensors",
+				{
+					headers: {
+						Range: "bytes=10000000-10099999",
+					},
+				}
+			).then((res) => res.arrayBuffer());
+
+			console.log("xet", xetDownload.byteLength, "bridge", bridgeDownload.byteLength);
+			expect(new Uint8Array(xetDownload)).toEqual(new Uint8Array(bridgeDownload));
+		}, 30_000);
+	}
+
+	it("should load text correctly when offset_into_range starts in a chunk further than the first", async () => {
+		const blob = new XetBlob({
+			repo: {
+				type: "model",
+				name: "celinah/xet-experiments",
+			},
+			hash: "794efea76d8cb372bbe1385d9e51c3384555f3281e629903ecb6abeff7d54eec",
+			size: 62_914_580,
+		});
+
+		// Reconstruction info
+		// {
+		// 	"offset_into_first_range": 600000,
+		// 	"terms":
+		// 		[
+		// 			{
+		// 				"hash": "be748f77930d5929cabd510a15f2c30f2f460b639804ef79dea46affa04fd8b2",
+		// 				"unpacked_length": 655360,
+		// 				"range": { "start": 0, "end": 5 },
+		// 			},
+		// 			{
+		// 				"hash": "be748f77930d5929cabd510a15f2c30f2f460b639804ef79dea46affa04fd8b2",
+		// 				"unpacked_length": 655360,
+		// 				"range": { "start": 0, "end": 5 },
+		// 			},
+		// 		],
+		// 	"fetch_info":
+		// 		{
+		// 			"be748f77930d5929cabd510a15f2c30f2f460b639804ef79dea46affa04fd8b2":
+		// 				[
+		// 					{
+		// 						"range": { "start": 0, "end": 5 },
+		// 						"url": "...",
+		// 						"url_range": { "start": 0, "end": 2839 },
+		// 					},
+		// 				],
+		// 		},
+		// }
+
+		const text = await blob.slice(600_000, 700_000).text();
+		const bridgeDownload = await fetch("https://huggingface.co/celinah/xet-experiments/resolve/main/large_text.txt", {
+			headers: {
+				Range: "bytes=600000-699999",
+			},
+		}).then((res) => res.text());
+
+		console.log("xet", text.length, "bridge", bridgeDownload.length);
+		expect(text.length).toBe(bridgeDownload.length);
+	}, 30_000);
+
+	describe("bg4_regoup_bytes", () => {
+		it("should regroup bytes when the array is %4 length", () => {
+			expect(bg4_regoup_bytes(new Uint8Array([1, 5, 2, 6, 3, 7, 4, 8]))).toEqual(
+				new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8])
+			);
+		});
+
+		it("should regroup bytes when the array is %4 + 1 length", () => {
+			expect(bg4_regoup_bytes(new Uint8Array([1, 5, 9, 2, 6, 3, 7, 4, 8]))).toEqual(
+				new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9])
+			);
+		});
+
+		it("should regroup bytes when the array is %4 + 2 length", () => {
+			expect(bg4_regoup_bytes(new Uint8Array([1, 5, 9, 2, 6, 10, 3, 7, 4, 8]))).toEqual(
+				new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+			);
+		});
+
+		it("should regroup bytes when the array is %4 + 3 length", () => {
+			expect(bg4_regoup_bytes(new Uint8Array([1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8]))).toEqual(
+				new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+			);
+		});
+	});
+});