@genai-fi/nanogpt 0.22.0 → 0.23.0

package/dist/training/tasks/TokenStore.js

@@ -0,0 +1,215 @@
+//#region lib/training/tasks/TokenStore.ts
+var e = "llm-tokenstore", t = class {
+	name;
+	tokeniserId;
+	datasetId;
+	shards = [];
+	masks;
+	shardCount = 0;
+	lastShardLength = -1;
+	_shardSize = 4e3 * 1024;
+	dirHandle = null;
+	opfsAvailable = !1;
+	lru = /* @__PURE__ */ new Map();
+	maxCachedShards;
+	constructor(t, n, r, i = 8, a = 4e3 * 1024) {
+		this.tokeniserId = t, this.datasetId = n, this.maxCachedShards = Math.max(1, Math.trunc(i)), this._shardSize = Math.max(1, Math.trunc(a)), this.name = r ?? e;
+	}
+	get shardSize() {
+		return this._shardSize;
+	}
+	hasMask() {
+		return !!this.masks;
+	}
+	getShardLength(e) {
+		if (e < 0 || e >= this.shardCount) throw RangeError(`Shard index ${e} out of range`);
+		return e === this.shardCount - 1 ? this.lastShardLength : this._shardSize;
+	}
+	async slice(e, t) {
+		if (e < 0 || t > this.getTokenCount() || e >= t) throw RangeError(`Invalid slice range: ${e} to ${t}`);
+		let n = new Uint16Array(t - e), r = 0, i = Math.floor(e / this._shardSize), a = Math.floor((t - 1) / this._shardSize);
+		for (let o = i; o <= a; o++) {
+			let i = o * this._shardSize, a = i + this.getShardLength(o), s = Math.max(e, i), c = Math.min(t, a);
+			if (s < c) {
+				let e = await this.getShard(o), t = s - i, a = c - i;
+				n.set(e.subarray(t, a), r), r += a - t;
+			}
+		}
+		return n;
+	}
+	async init() {
+		try {
+			let e = globalThis.navigator;
+			if (!e?.storage?.getDirectory) {
+				this.opfsAvailable = !1;
+				return;
+			}
+			let t = await e.storage.getDirectory();
+			this.dirHandle = await t.getDirectoryHandle(this.name, { create: !0 }), this.opfsAvailable = !0;
+			try {
+				let e = await this.readManifest();
+				if (e && (e?.tokenizerId !== this.tokeniserId || e?.datasetId !== this.datasetId)) for await (let e of this.dirHandle.values()) {
+					let t = e.name;
+					if (t.startsWith("shard-") || t.startsWith("mask-") || t === "manifest.json") try {
+						await this.dirHandle.removeEntry(t);
+					} catch {}
+				}
+				else e && (this.shardCount = e.shardCount, this._shardSize = e.shardSize, this.lastShardLength = e.lastShardLength, this.shards = Array(this.shardCount).fill(null), e.maskExistsIndex && Array.isArray(e.maskExistsIndex) && (this.masks = Array(this.shardCount).fill(null)));
+			} catch {}
+		} catch {
+			this.opfsAvailable = !1, this.dirHandle = null, console.warn("TokenStore: OPFS initialization failed, falling back to memory-only storage");
+		}
+	}
+	async writeManifest() {
+		if (!this.opfsAvailable || !this.dirHandle) return;
+		let e = {
+			tokenizerId: this.tokeniserId,
+			datasetId: this.datasetId,
+			shardSize: this._shardSize,
+			shardCount: this.shardCount,
+			lastShardLength: this.lastShardLength,
+			maskExistsIndex: (this.masks || []).map((e) => !!e),
+			timestamp: Date.now()
+		}, t = await (await this.dirHandle.getFileHandle("manifest.json", { create: !0 })).createWritable({ keepExistingData: !1 });
+		await t.write(JSON.stringify(e)), await t.close();
+	}
+	async readManifest() {
+		if (!(!this.opfsAvailable || !this.dirHandle)) try {
+			let e = await (await (await this.dirHandle.getFileHandle("manifest.json")).getFile()).text();
+			return JSON.parse(e);
+		} catch {
+			return;
+		}
+	}
+	shardFilename(e) {
+		return `shard-${e}.bin`;
+	}
+	maskFilename(e) {
+		return `mask-${e}.bin`;
+	}
+	async writeShardToOPFS(e, t) {
+		if (!this.opfsAvailable || !this.dirHandle) return;
+		let n = this.shardFilename(e), r = await (await this.dirHandle.getFileHandle(n, { create: !0 })).createWritable();
+		await r.write(t.buffer), await r.close();
+	}
+	async writeMaskToOPFS(e, t) {
+		if (!this.opfsAvailable || !this.dirHandle) return;
+		let n = this.maskFilename(e), r = await (await this.dirHandle.getFileHandle(n, { create: !0 })).createWritable();
+		await r.write(t.buffer), await r.close();
+	}
+	async readShardFromOPFS(e) {
+		if (!(!this.opfsAvailable || !this.dirHandle)) try {
+			let t = await (await (await this.dirHandle.getFileHandle(this.shardFilename(e))).getFile()).arrayBuffer();
+			return new Uint16Array(t);
+		} catch {
+			return;
+		}
+	}
+	async readMaskFromOPFS(e) {
+		if (!(!this.opfsAvailable || !this.dirHandle)) try {
+			let t = await (await (await this.dirHandle.getFileHandle(this.maskFilename(e))).getFile()).arrayBuffer();
+			return new Uint8Array(t);
+		} catch {
+			return;
+		}
+	}
+	touchLRU(e) {
+		this.lru.delete(e), this.lru.set(e, !0);
+	}
+	evictIfNeeded() {
+		if (!(this.maxCachedShards <= 0)) for (; this.lru.size > this.maxCachedShards;) {
+			let e = this.lru.keys().next().value;
+			if (e === void 0) break;
+			this.lru.delete(e), this.shards[e] && (this.shards[e] = null), this.masks?.[e] && (this.masks[e] = null);
+		}
+	}
+	async getShard(e) {
+		if (!Number.isInteger(e) || e < 0 || e >= this.shards.length) throw RangeError(`Shard index ${e} out of range`);
+		let t = this.shards[e];
+		if (t) return this.touchLRU(e), t;
+		if (this.opfsAvailable) {
+			let t = await this.readShardFromOPFS(e);
+			if (t) return this.shards[e] = t, this.touchLRU(e), this.evictIfNeeded(), t;
+		}
+		throw Error(`Shard ${e} not available in memory or OPFS`);
+	}
+	async getMask(e) {
+		if (!Number.isInteger(e) || e < 0 || e >= this.shardCount || !this.masks) return;
+		let t = this.masks[e];
+		if (t) return this.touchLRU(e), t;
+		if (this.opfsAvailable) {
+			let t = await this.readMaskFromOPFS(e);
+			if (t) return this.masks[e] = t, this.touchLRU(e), this.evictIfNeeded(), t;
+		}
+	}
+	getShardCount() {
+		return this.shardCount;
+	}
+	getTokenCount() {
+		return this.shardCount === 0 ? 0 : (this.shardCount - 1) * this._shardSize + this.lastShardLength;
+	}
+	async appendShard(e, t) {
+		if (this.lastShardLength >= 0 && this.lastShardLength < this._shardSize) throw Error("Previous shard was not full");
+		let n = this.shards.length;
+		this.shards.push(e), !this.masks && t && (this.masks = Array(n).fill(null)), this.masks && this.masks.push(t ?? null), this.shardCount = this.shards.length, this.lastShardLength = e.length, this.touchLRU(n);
+		try {
+			this.opfsAvailable && this.dirHandle && (await this.writeShardToOPFS(n, e), t && await this.writeMaskToOPFS(n, t), await this.writeManifest());
+		} catch (e) {
+			console.error(e), this.maxCachedShards = 2 ** 53 - 1;
+		}
+		this.evictIfNeeded();
+	}
+	async dispose() {
+		if (this.opfsAvailable && this.dirHandle) try {
+			for await (let e of this.dirHandle.values()) {
+				let t = e.name;
+				try {
+					await this.dirHandle.removeEntry(t);
+				} catch {}
+			}
+			try {
+				let e = globalThis.navigator;
+				if (e?.storage?.getDirectory) {
+					let t = await e.storage.getDirectory();
+					t.removeEntry && await t.removeEntry(this.name, { recursive: !0 });
+				}
+			} catch {}
+		} catch {}
+		this.shards = [], this.masks = [], this.shardCount = 0, this.lastShardLength = -1, this.lru.clear(), this.dirHandle = null, this.opfsAvailable = !1;
+	}
+	async clear() {
+		if (this.shards = [], this.masks = [], this.shardCount = 0, this.lastShardLength = -1, this.lru.clear(), this.opfsAvailable && this.dirHandle) try {
+			for await (let e of this.dirHandle.values()) {
+				let t = e.name;
+				try {
+					await this.dirHandle.removeEntry(t);
+				} catch {}
+			}
+			await this.writeManifest();
+		} catch {}
+	}
+}, n = /* @__PURE__ */ new Map();
+function r(e, t, r) {
+	let i = n.get(e);
+	return i && i.tokeniserId === t && i.datasetId === r ? i : null;
+}
+async function i(e) {
+	let t = n.get(e);
+	if (t) await t.dispose(), n.delete(e);
+	else try {
+		let t = globalThis.navigator;
+		if (t?.storage?.getDirectory) {
+			let n = await t.storage.getDirectory();
+			n.removeEntry && await n.removeEntry(e, { recursive: !0 });
+		}
+	} catch {}
+}
+async function a(e, r, i, a) {
+	let o = n.get(e);
+	if (o && o.tokeniserId === r && o.datasetId === i) return o;
+	o && (await o.dispose(), n.delete(e));
+	let s = new t(r, i, e, a?.maxCachedShards ?? 8, a?.shardSize ?? 4e3 * 1024);
+	return a?.noOPFS || await s.init(), n.set(e, s), s;
+}
+//#endregion
+export { t as TokenStore, a as createTokenStore, i as deleteTokenStore, r as getTokenStore };

package/dist/training/tasks/splitter.js

@@ -1,4 +1,4 @@
-import { n as e } from "../../stream-DKl3GTDL.js";
+import { n as e } from "../../stream-BjdpSNqB.js";
 import t from "./ConversationTask.js";
 //#region lib/training/tasks/splitter.ts
 async function n(n, r) {

package/dist/training/validation.d.ts

@@ -1,8 +1,10 @@
-import { ITokeniser, Task } from '../../main';
+import { ITokeniser } from '../../tokeniser/type';
 import { Tensor } from '@tensorflow/tfjs-core';
 import { Dataset } from '@tensorflow/tfjs-data';
 import { DatasetBuilder, DatasetState } from './DatasetBuilder';
-export declare function createTrainValidationSplit(tasks: Task[] | Uint16Array[], tokeniser: ITokeniser, datasetBuilder: DatasetBuilder, batchSize: number, validationSplit?: number, masking?: boolean): Promise<{
+import { TokenStore } from './tasks/TokenStore';
+export declare function storeFromArray(tokens: Uint16Array[], tokenizer: ITokeniser): Promise<TokenStore>;
+export declare function createTrainValidationDatasets(trainingTokens: Uint16Array[] | TokenStore, validationTokens: Uint16Array[] | TokenStore, tokeniser: ITokeniser, datasetBuilder: DatasetBuilder, batchSize: number): Promise<{
     trainDataset: Dataset<{
         xs: Tensor;
         ys: Tensor;

package/dist/training/validation.js

@@ -1,2 +1,21 @@
-import { o as e } from "../main-Bgc7_9kb.js";
-export { e as createTrainValidationSplit };
+import { TokenStore as e, createTokenStore as t } from "./tasks/TokenStore.js";
+//#region lib/training/validation.ts
+async function n(e, n) {
+	let r = await t("training-tokens", n.id, n.datasetID ?? ""), i = e.map((e) => r.appendShard(e));
+	return await Promise.all(i), r;
+}
+async function r(t, r, i, a, o) {
+	let s = t instanceof e ? t : await n(t, i), c = r instanceof e ? r : await n(r, i), l = s.getTokenCount(), { dataset: u, state: d } = await a.createTextDataset(s, { batchSize: o }), { dataset: f, state: p } = await a.createTextDataset(c, {
+		batchSize: o,
+		shuffleFirst: !0
+	});
+	return {
+		trainDataset: u,
+		validationDataset: f,
+		size: l,
+		validationState: p,
+		trainState: d
+	};
+}
+//#endregion
+export { r as createTrainValidationDatasets, n as storeFromArray };

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.22.0",
+    "version": "0.23.0",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",