@genai-fi/nanogpt 0.21.0 → 0.23.0

package/dist/training/tasks/Task.js

@@ -1,9 +1,10 @@
 import { yieldIfNeeded as e } from "../../utilities/yielder.js";
+import { createTokenStore as t, deleteTokenStore as n } from "./TokenStore.js";
 //#region lib/training/tasks/Task.ts
-var t = class {};
-function n(e, t, n, r, i, a) {
+var r = class {};
+async function i(e, t, n, r, i, a) {
 	for (let o of e) {
-		let e = o.nextTokens(n, a ? !0 : void 0);
+		let e = await o.nextTokens(n, a ? !0 : void 0);
 		if (e) {
 			let n = Array.isArray(e) ? e : e.tokens;
 			r.total += n.length;
@@ -11,7 +12,9 @@ function n(e, t, n, r, i, a) {
 			if (r.offset + n.length > o.length) {
 				let c = o.length - r.offset;
 				o.set(n.slice(0, c), r.offset);
-				let l = n.length - c, u = new Uint16Array(Math.max(Math.floor(i * .1) + 100, l));
+				let l = n.length - c;
+				if (l > i) throw Error(`Estimated tokens (${i}) is too small for the next batch of tokens (${l}).`);
+				let u = new Uint16Array(i);
 				if (u.set(n.slice(c), 0), t.push(u), a && s && !Array.isArray(e)) {
 					s.set(e.mask.slice(0, c).map((e) => +!!e), r.offset);
 					let t = new Uint8Array(u.length);
@@ -22,34 +25,26 @@ function n(e, t, n, r, i, a) {
 		}
 	}
 }
-async function r(t, r, i, a) {
-	let o = Math.min((await Promise.all(t.map((e) => e.estimateTokens(r)))).reduce((e, t) => e + t, 0), r.vocabSize * 1e4), s = [new Uint16Array(o)], c = a ? [new Uint8Array(o)] : null, l = {
+async function a(r, a, o) {
+	await n("training-tokens");
+	let s = await t("training-tokens", a.id, a.datasetID ?? "", o);
+	await n("validation-tokens");
+	let c = o?.validationSplit && o.validationSplit > 0 ? await t("validation-tokens", a.id, a.datasetID ?? "", o) : void 0, l = [new Uint16Array(s.shardSize)], u = o?.masking ? [new Uint8Array(s.shardSize)] : null, d = {
 		offset: 0,
 		total: 0
-	}, u = performance.now();
-	for (; l.offset < o && (n(t, s, r, l, o, c || void 0), !t.every((e) => !e.hasMoreConversations()));) u = await e(u, i, l.total);
-	if (s.length === 1) return c ? {
-		tokens: s[0].subarray(0, l.offset),
-		mask: c[0].subarray(0, l.offset)
-	} : s[0].subarray(0, l.offset);
-	let d = s.reduce((e, t) => e + t.length, 0) - (s[s.length - 1].length - l.offset), f = new Uint16Array(d), p = 0;
-	for (let e = 0; e < s.length; e++) {
-		let t = s[e];
-		e === s.length - 1 ? (f.set(t.subarray(0, l.offset), p), p += l.offset) : (f.set(t, p), p += t.length);
-	}
-	if (c) {
-		let e = new Uint8Array(d);
-		p = 0;
-		for (let t = 0; t < c.length; t++) {
-			let n = c[t];
-			t === c.length - 1 ? (e.set(n.subarray(0, l.offset), p), p += l.offset) : (e.set(n, p), p += n.length);
-		}
-		return {
-			tokens: f,
-			mask: e
-		};
+	}, f = o?.validationSplit && o.validationSplit > 0 ? [new Uint16Array(c.shardSize)] : void 0, p = o?.masking && f ? [new Uint8Array(c.shardSize)] : null, m = {
+		offset: 0,
+		total: 0
+	}, h = performance.now();
+	for (;;) {
+		let t = o?.validationSplit && o.validationSplit > 0 && Math.random() < o.validationSplit, n = t ? f : l, g = t ? m : d, _ = t ? p : u, v = t ? c : s;
+		if (await i(r, n, a, g, v.shardSize, _ || void 0), n.length > 1 && (v.appendShard(n[0], _ ? _[0] : void 0), n.shift(), _ && _.shift()), r.every((e) => !e.hasMoreConversations())) break;
+		h = await e(h, o?.cb, d.total);
 	}
-	return f;
+	return l.length === 1 && (l[0] = l[0].subarray(0, d.offset), await s.appendShard(l[0], u ? u[0].subarray(0, d.offset) : void 0)), f && f.length === 1 && (f[0] = f[0].subarray(0, m.offset), await c.appendShard(f[0], p ? p[0].subarray(0, m.offset) : void 0)), {
+		trainingTokens: s,
+		validationTokens: f ? c : void 0
+	};
 }
 //#endregion
-export { t as Task, r as tokensFromTasks };
+export { r as Task, a as tokensFromTasks };

package/dist/training/tasks/TokenStore.d.ts

@@ -0,0 +1,46 @@
+export declare class TokenStore {
+    readonly name: string;
+    readonly tokeniserId: string;
+    readonly datasetId: string;
+    private shards;
+    private masks?;
+    private shardCount;
+    private lastShardLength;
+    private _shardSize;
+    private dirHandle;
+    private opfsAvailable;
+    private lru;
+    private maxCachedShards;
+    constructor(tokeniserId: string, datasetId: string, name?: string, maxCachedShards?: number, shardSize?: number);
+    get shardSize(): number;
+    hasMask(): boolean;
+    getShardLength(index: number): number;
+    slice(start: number, end: number): Promise<Uint16Array>;
+    init(): Promise<void>;
+    private writeManifest;
+    private readManifest;
+    private shardFilename;
+    private maskFilename;
+    private writeShardToOPFS;
+    private writeMaskToOPFS;
+    private readShardFromOPFS;
+    private readMaskFromOPFS;
+    private touchLRU;
+    private evictIfNeeded;
+    getShard(index: number): Promise<Uint16Array>;
+    getMask(index: number): Promise<Uint8Array | undefined>;
+    getShardCount(): number;
+    getTokenCount(): number;
+    appendShard(shard: Uint16Array, mask?: Uint8Array): Promise<void>;
+    dispose(): Promise<void>;
+    clear(): Promise<void>;
+}
+export declare function getTokenStore(name: string, tokeniserId: string, datasetId: string): TokenStore | null;
+export declare function deleteTokenStore(name: string): Promise<void>;
+interface TokenStoreOptions {
+    maxCachedShards?: number;
+    shardSize?: number;
+    noOPFS?: boolean;
+}
+export declare function createTokenStore(name: string, tokeniserId: string, datasetId: string, options?: TokenStoreOptions): Promise<TokenStore>;
+export {};

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.d.ts

@@ -1,5 +1,5 @@
 import { Task } from './Task';
-export default function splitValidation(tasks: Task[], validationSplit: number): {
+export default function splitValidation(tasks: Task[], validationSplit: number): Promise<{
     training: Task;
     validation: Task;
-};
+}>;

package/dist/training/tasks/splitter.js

@@ -1,18 +1,18 @@
-import e from "./ConversationTask.js";
+import { n as e } from "../../stream-BjdpSNqB.js";
+import t from "./ConversationTask.js";
 //#region lib/training/tasks/splitter.ts
-function t(t, n) {
-	if (n <= 0 || n >= 1) throw Error("validationSplit must be between 0 and 1");
-	t.forEach((e) => e.shuffle());
-	let r = [], i = [];
-	for (let e of t) for (; e.hasMoreConversations();) {
-		let t = e.nextConversation();
+async function n(n, r) {
+	if (r <= 0 || r >= 1) throw Error("validationSplit must be between 0 and 1");
+	let i = [], a = [];
+	for (let e of n) for (; e.hasMoreConversations();) {
+		let t = await e.nextConversation();
 		if (!t) break;
-		Math.random() < n ? i.push(t) : r.push(t);
+		Math.random() < r ? a.push(t) : i.push(t);
 	}
 	return {
-		training: new e(r),
-		validation: new e(i)
+		training: new t([new e(i)]),
+		validation: new t([new e(a)])
 	};
 }
 //#endregion
-export { t as default };
+export { n as default };

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-DdsoTJAW.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/dist/utilities/tokens.d.ts

@@ -0,0 +1,2 @@
+export declare function sliceUint16Shards(shards: readonly Uint16Array[], start: number, endExclusive: number, shardSize?: number): Uint16Array;
+export declare function sliceUint8Shards(shards: readonly Uint8Array[], start: number, endExclusive: number, shardSize?: number): Uint8Array;

package/dist/utilities/tokens.js

@@ -0,0 +1,40 @@
+//#region lib/utilities/tokens.ts
+function e(e, t, n, r = e[0]?.length ?? 0) {
+	if (e.length === 0) return new Uint16Array();
+	if (!Number.isInteger(t) || !Number.isInteger(n)) throw RangeError("start/endExclusive must be integers");
+	if (t < 0 || n < t) throw RangeError(`Invalid range [${t}, ${n})`);
+	if (r <= 0) throw RangeError("Invalid shardSize");
+	let i = e[e.length - 1], a = (e.length - 1) * r + i.length;
+	if (n > a) throw RangeError(`Slice end ${n} is out of bounds for total length ${a}`);
+	if (t === n) return new Uint16Array();
+	let o = new Uint16Array(n - t), s = 0, c = t;
+	for (; c < n;) {
+		let t = Math.floor(c / r), i = c - t * r, a = e[t], l = Math.min(n - c, a.length - i);
+		try {
+			o.set(a.subarray(i, i + l), s);
+		} catch (n) {
+			throw console.error(`Error slicing shards: shardIndex=${t}, shardOffset=${i}, take=${l}, writeOffset=${s}`), e.forEach((t, n) => {
+				t.length !== r && n !== e.length - 1 && console.error(`Shard ${n} has length ${t.length}, expected ${r}`);
+			}), n;
+		}
+		c += l, s += l;
+	}
+	return o;
+}
+function t(e, t, n, r = e[0]?.length ?? 0) {
+	if (e.length === 0) return new Uint8Array();
+	if (!Number.isInteger(t) || !Number.isInteger(n)) throw RangeError("start/endExclusive must be integers");
+	if (t < 0 || n < t) throw RangeError(`Invalid range [${t}, ${n})`);
+	if (r <= 0) throw RangeError("Invalid shardSize");
+	let i = e[e.length - 1], a = (e.length - 1) * r + i.length;
+	if (n > a) throw RangeError(`Slice end ${n} is out of bounds for total length ${a}`);
+	if (t === n) return new Uint8Array();
+	let o = new Uint8Array(n - t), s = 0, c = t;
+	for (; c < n;) {
+		let t = Math.floor(c / r), i = c - t * r, a = e[t], l = Math.min(n - c, a.length - i);
+		o.set(a.subarray(i, i + l), s), c += l, s += l;
+	}
+	return o;
+}
+//#endregion
+export { e as sliceUint16Shards, t as sliceUint8Shards };

package/dist/utilities/topP.js

@@ -6,6 +6,12 @@ function e(e, t) {
 	})).sort((e, t) => t.prob - e.prob), r = 0, i = Array(n.length).fill(0);
 	for (let e of n) if (r += e.prob, i[e.index] = e.prob, r >= t) break;
 	let a = i.reduce((e, t) => e + t, 0);
+	if (a === 0) {
+		let t = e[0], n = t.reduce((e, t) => e + t, 0);
+		if (n > 0) return t.map((e) => e / n);
+		let r = 1 / t.length;
+		return t.map(() => r);
+	}
 	return i.map((e) => e / a);
 }
 //#endregion

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.21.0",
+    "version": "0.23.0",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",
@@ -52,6 +52,7 @@
     "dependencies": {
         "@tensorflow/tfjs": "^4.22.0",
         "@tensorflow/tfjs-backend-webgpu": "^4.22.0",
+        "@zip.js/zip.js": "^2.8.34",
         "eventemitter3": "^5.0.1",
         "jszip": "^3.10.1",
         "papaparse": "^5.5.3",