@genai-fi/nanogpt 0.22.0 → 0.24.0

package/dist/tokeniser/BaseTokeniser.js

@@ -1,2 +1,135 @@
-import { n as e, t } from "../BaseTokeniser-C9TSv4th.js";
-export { e as SPECIALS, t as default };
+import { t as e } from "../eventemitter3-D_qV3Lof.js";
+//#region lib/tokeniser/BaseTokeniser.ts
+var t = [
+	"<eos>",
+	"<bos>",
+	"",
+	"<pad>",
+	"<|user_start|>",
+	"<|user_end|>",
+	"<|assistant_start|>",
+	"<|assistant_end|>",
+	"<|system_start|>",
+	"<|system_end|>"
+], n = class extends e {
+	id = "untrained";
+	datasetID;
+	specialTokens = /* @__PURE__ */ new Map();
+	specialTokenSet = /* @__PURE__ */ new Set();
+	isSpecialToken(e) {
+		return this.specialTokenSet.has(e);
+	}
+	addSpecialTokens() {
+		t.forEach((e, t) => {
+			this.addToken(e, t), this.specialTokens.set(e, t), this.specialTokenSet.add(t);
+		});
+	}
+	addSpecialToken(e, t) {
+		this.specialTokens.set(e, t), this.specialTokenSet.add(t);
+	}
+	generateID() {
+		let e = this.getVocab(), t = 2166136261, n = 2654435769;
+		if (e.length === 0) {
+			this.id = "untrained";
+			return;
+		}
+		for (let r = 0; r < e.length; r++) {
+			let i = e[r];
+			t ^= i.length, t = Math.imul(t, 16777619), n ^= r, n = Math.imul(n, 2246822507);
+			for (let e = 0; e < i.length; e++) {
+				let r = i.charCodeAt(e);
+				t ^= r, t = Math.imul(t, 16777619), n ^= r, n = Math.imul(n, 3266489909);
+			}
+		}
+		let r = (t >>> 0).toString(36), i = (n >>> 0).toString(36);
+		this.id = "tokeniser_" + r + "_" + i;
+	}
+	encodeSequence(e) {
+		let t = this.encode(e);
+		return [
+			this.bosToken,
+			...t,
+			this.eosToken
+		];
+	}
+	encodeAsSequence(e, t) {
+		let n = e.flatMap((e) => this.encode(e.content));
+		return t ? [
+			this.bosToken,
+			...n,
+			this.eosToken,
+			this.bosToken
+		] : [
+			this.bosToken,
+			...n,
+			this.eosToken
+		];
+	}
+	encodeConversation(e, t, n) {
+		let r = [[this.bosToken]], i;
+		n && (i = [[!1]]);
+		let a = [
+			this.getSpecialTokenIndex("<|user_start|>"),
+			this.getSpecialTokenIndex("<|assistant_start|>"),
+			this.getSpecialTokenIndex("<|system_start|>")
+		], o = [
+			this.getSpecialTokenIndex("<|user_end|>"),
+			this.getSpecialTokenIndex("<|assistant_end|>"),
+			this.getSpecialTokenIndex("<|system_end|>")
+		];
+		for (let t of e) {
+			let e = !1, s = this.encode(t.content);
+			switch (t.role) {
+				case "user":
+					r.push([a[0]]), e = !0;
+					break;
+				case "assistant":
+					r.push([a[1]]);
+					break;
+				case "system":
+					r.push([a[2]]), e = !0;
+					break;
+			}
+			switch (r.push(s), t.role) {
+				case "user":
+					r.push([o[0]]);
+					break;
+				case "assistant":
+					r.push([o[1]]);
+					break;
+				case "system":
+					r.push([o[2]]);
+					break;
+			}
+			n && i && e ? (i.push([!1]), i.push(s.map(() => !1)), i.push([!1])) : n && i && (i.push([!1]), i.push(s.map(() => !0)), i.push([!0]));
+		}
+		let s = r.flat();
+		return t ? (s.push(a[1]), n && i && i.push([!1])) : (s.push(this.eosToken), n && i && i.push([!0])), n && i ? {
+			tokens: s,
+			mask: i.flat()
+		} : s;
+	}
+	decodeConversation(e) {
+		let t = [], n = 0;
+		for (; n < e.length;) {
+			let r = e[n], i = null;
+			if (r === this.getSpecialTokenIndex("<|user_start|>") ? i = "user" : r === this.getSpecialTokenIndex("<|assistant_start|>") ? i = "assistant" : r === this.getSpecialTokenIndex("<|system_start|>") ? i = "system" : r === this.bosToken || (r === this.eosToken ? i = null : (i = "text", n--)), i) {
+				n++;
+				let r = [];
+				for (; n < e.length && e[n] !== this.getSpecialTokenIndex(`<|${i}_end|>`) && e[n] !== this.eosToken;) r.push(e[n]), n++;
+				let a = this.decode(r);
+				t.push({
+					role: i,
+					content: a
+				});
+			}
+			n++;
+		}
+		return t;
+	}
+	getSpecialTokenIndex(e) {
+		return this.specialTokens.get(e);
+	}
+};
+//#endregion
+export { t as SPECIALS, n as default };

package/dist/tokeniser/CharTokeniser.js

@@ -1,7 +1,6 @@
-import { yieldIfNeeded as e } from "../utilities/yielder.js";
-import { n as t, t as n } from "../BaseTokeniser-C9TSv4th.js";
+import e, { SPECIALS as t } from "./BaseTokeniser.js";
 //#region lib/tokeniser/CharTokeniser.ts
-var r = ["<eos>", "<unk>"], i = class extends n {
+var n = ["<eos>", "<unk>"], r = class extends e {
 	vocabSize = 0;
 	eosToken = 0;
 	bosToken = 0;
@@ -35,22 +34,21 @@ var r = ["<eos>", "<unk>"], i = class extends n {
 	destroy() {
 		this.cache.clear(), this.vocab = [];
 	}
-	async train(t, n, i) {
-		this.datasetID = i;
-		let a = /* @__PURE__ */ new Set(), o = performance.now();
-		for (let r of t) {
-			let t = r.cursor(), i = await t.next();
-			for (; i !== null;) i.forEach((e) => {
-				for (let t of e.content) a.add(t);
-			}), o = await e(o, n, 0), i = await t.next();
-		}
-		let s = Array.from(a), c = this.vocab.indexOf("", this.unkToken + 1), l = this.vocabSize - r.length;
-		if (c === -1) return this.generateID(), this.vocabSize;
-		if (this._trained = !0, s.length > l) throw Error("too_small_vocab");
-		let u = c;
-		if (u !== -1) {
+	async train(e, t, r) {
+		this.datasetID = r;
+		let i = /* @__PURE__ */ new Set();
+		for (let n of e) await n.begin((e) => {
+			e.forEach((e) => {
+				for (let t of e.content) i.add(t);
+			});
+		}, t ? () => t(i.size) : void 0);
+		let a = Array.from(i), o = this.vocab.indexOf("", this.unkToken + 1), s = this.vocabSize - n.length;
+		if (o === -1) return this.generateID(), this.vocabSize;
+		if (this._trained = !0, a.length > s) throw Error("too_small_vocab");
+		let c = o;
+		if (c !== -1) {
 			let e = new Set(this.vocab);
-			for (let t of s) if (!e.has(t) && (this.vocab[u] = t, e.add(t), u = this.vocab.indexOf("", u + 1), u === -1)) break;
+			for (let t of a) if (!e.has(t) && (this.vocab[c] = t, e.add(t), c = this.vocab.indexOf("", c + 1), c === -1)) break;
 		}
 		return this.cache.clear(), this.vocab.forEach((e, t) => {
 			this.cache.set(e, t);
@@ -85,4 +83,4 @@ var r = ["<eos>", "<unk>"], i = class extends n {
 	}
 };
 //#endregion
-export { i as default };
+export { r as default };

package/dist/tokeniser/bpe.js

@@ -1,5 +1,5 @@
-import { yieldIfNeeded as e } from "../utilities/yielder.js";
-import { n as t, t as n } from "../BaseTokeniser-C9TSv4th.js";
+import e, { SPECIALS as t } from "./BaseTokeniser.js";
+import { yieldIfNeeded as n } from "../utilities/yielder.js";
 import r from "../utilities/tokenParse.js";
 //#region lib/tokeniser/bpe.ts
 function i(e, t) {
@@ -58,7 +58,7 @@ function l(e, t) {
 		e.tokens[n] = i;
 	}), e.pairs.delete(i(t.a, t.b));
 }
-var u = class extends n {
+var u = class extends e {
 	targetSize;
 	vocab = /* @__PURE__ */ new Set();
 	vocabIndex = /* @__PURE__ */ new Map();
@@ -98,26 +98,22 @@ var u = class extends n {
 	get unkToken() {
 		return this.vocabIndex.get("") ?? 1;
 	}
-	async train(t = [], n, i) {
+	async train(e = [], t, i) {
 		this.datasetID = i;
-		let o = performance.now(), c = /* @__PURE__ */ new Set();
-		for (let i of t) {
-			let t = i.cursor(), a = await t.next();
-			for (; a !== null;) {
-				for (let e of a) {
-					let t = r(e.content);
-					for (let e of t) c.add(e);
-				}
-				o = await e(o, n, this.vocab.size), a = await t.next();
-			}
-		}
+		let o = /* @__PURE__ */ new Set();
 		this.vocab = /* @__PURE__ */ new Set(), this.pretokenMap.clear(), this.merges = [], this.addSpecialTokens();
-		let u = Array.from(c), d = u.map((e) => Array.from(e).map((e) => (this.vocab.add(e), e))), f = a(d);
-		if (o = await e(o, n, this.vocab.size), this.vocab.size >= this.targetSize) throw console.warn("Initial vocab size is greater than or equal to target size. No merges will be performed.", this.vocab.size, this.targetSize), Error("too_small_vocab");
+		for (let n of e) await n.begin((e) => {
+			for (let t of e) {
+				let e = r(t.content);
+				for (let t of e) o.has(t) || (o.add(t), Array.from(t).forEach((e) => this.vocab.add(e)));
+			}
+		}, t ? () => t(this.vocab.size) : void 0);
+		let c = performance.now(), u = Array.from(o), d = u.map((e) => Array.from(e)), f = a(d);
+		if (c = await n(c, t, this.vocab.size), this.vocab.size >= this.targetSize) throw console.warn("Initial vocab size is greater than or equal to target size. No merges will be performed.", this.vocab.size, this.targetSize), Error("too_small_vocab");
 		for (; this.vocab.size < this.targetSize && this.merges.length < this.targetSize;) {
-			let t = s(f);
-			if (!t) break;
-			this.merges.push([t.a, t.b]), this.vocab.add(t.a + t.b), l(f, t), o = await e(o, n, this.vocab.size);
+			let e = s(f);
+			if (!e) break;
+			this.merges.push([e.a, e.b]), this.vocab.add(e.a + e.b), l(f, e), c = await n(c, t, this.vocab.size);
 		}
 		u.forEach((e, t) => {
 			let n = d[t];

package/dist/training/DatasetBuilder.d.ts

@@ -1,21 +1,36 @@
 import { Tensor } from '@tensorflow/tfjs-core';
 import { Conversation, ITokeniser } from '../tokeniser/type';
 import { Dataset } from '@tensorflow/tfjs-data';
+import { TokenStore } from './tasks/TokenStore';
 export declare function flattenTokens(textData: Conversation[][], tokenizer: ITokeniser): Uint16Array;
 export declare function flattenTokensWithMask(textData: Conversation[][], tokenizer: ITokeniser): {
     tokens: Uint16Array;
     mask: Uint8Array;
 };
-export declare function shuffle(array: Uint32Array): Uint32Array;
+export declare function shuffle(array: Uint32Array | Uint16Array): Uint32Array | Uint16Array;
 export interface DatasetState {
-    shuffledIndexes: Uint32Array;
+    shuffledShards: Uint16Array;
+    shuffledIndexes: Uint16Array;
+    lastShardIndexes: Uint16Array;
+    currentShard: Uint16Array | null;
+    nextShard: Uint16Array | null;
+    currentMask: Uint8Array | null;
+    nextMask: Uint8Array | null;
+    shardIndex: number;
     step: number;
 }
+export declare function moveToNext(state: DatasetState, store: TokenStore, noShuffle?: boolean): Promise<void>;
+interface DatasetOptions {
+    batchSize: number;
+    noShuffle?: boolean;
+    ignoreIndex?: number;
+    shuffleFirst?: boolean;
+}
 export declare class DatasetBuilder {
     tokenizer: ITokeniser;
     blockSize: number;
     constructor(tokenizer: ITokeniser, blockSize?: number);
-    createTextDataset(flatTokens: Uint16Array[], batchSize?: number, indexes?: Uint32Array, mask?: Uint8Array[], ignoreIndex?: number): Promise<{
+    createTextDataset(store: TokenStore, options?: DatasetOptions): Promise<{
         dataset: Dataset<{
             xs: Tensor;
             ys: Tensor;
@@ -23,3 +38,4 @@ export declare class DatasetBuilder {
         state: DatasetState;
     }>;
 }
+export {};

package/dist/training/DatasetBuilder.js

@@ -1,2 +1,2 @@
-import { i as e, n as t, r as n, t as r } from "../DatasetBuilder-Ctb425Id.js";
-export { r as DatasetBuilder, t as flattenTokens, n as flattenTokensWithMask, e as shuffle };
+import { a as e, i as t, n, r, t as i } from "../DatasetBuilder-DU1G1OKX.js";
+export { i as DatasetBuilder, n as flattenTokens, r as flattenTokensWithMask, t as moveToNext, e as shuffle };

package/dist/training/PreTrainer.js

@@ -1,4 +1,4 @@
-import { t as e } from "../DatasetBuilder-Ctb425Id.js";
+import { t as e } from "../DatasetBuilder-DU1G1OKX.js";
 import t from "./BasicTrainer.js";
 //#region lib/training/PreTrainer.ts
 var n = {

package/dist/training/SFTTrainer.js

@@ -1,4 +1,4 @@
-import { t as e } from "../DatasetBuilder-Ctb425Id.js";
+import { t as e } from "../DatasetBuilder-DU1G1OKX.js";
 import t from "./BasicTrainer.js";
 //#region lib/training/SFTTrainer.ts
 var n = {

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

@@ -0,0 +1,47 @@
+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;
+    finish(): Promise<void>;
+    appendShard(shard: Uint16Array, mask?: Uint8Array): 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,218 @@
+//#region lib/training/tasks/TokenStore.ts
+var e = "llm-tokenstore", t = class {
+	name;
+	tokeniserId;
+	datasetId;
+	shards = [];
+	masks;
+	shardCount = 0;
+	lastShardLength = -1;
+	_shardSize = 8e3 * 1024;
+	dirHandle = null;
+	opfsAvailable = !1;
+	lru = /* @__PURE__ */ new Map();
+	maxCachedShards;
+	constructor(t, n, r, i = 2, a = 8e3 * 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 finish() {
+		await this.writeManifest();
+	}
+	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 && (this.writeShardToOPFS(n, e), t && this.writeMaskToOPFS(n, t));
+		} 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, a?.shardSize);
+	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/tokenStream.d.ts

@@ -0,0 +1,16 @@
+import { ConversationStream, ITokeniser } from '../../../main';
+import { TokenStore } from './TokenStore';
+interface TokensFromTasksOptions {
+    masking?: boolean;
+    maxCachedShards?: number;
+    noOPFS?: boolean;
+    shardSize?: number;
+    validationSplit?: number;
+    validationSeed?: string | number;
+    cb?: (tokens: number) => void;
+}
+export declare function tokensFromStreams(tasks: ConversationStream[], tokenizer: ITokeniser, options?: TokensFromTasksOptions): Promise<{
+    trainingTokens: TokenStore;
+    validationTokens?: TokenStore;
+}>;
+export {};

package/dist/training/tasks/tokenStream.js

@@ -0,0 +1,46 @@
+import { createTokenStore as e, deleteTokenStore as t } from "./TokenStore.js";
+import { seededRng as n } from "../../utilities/random.js";
+//#region lib/training/tasks/tokenStream.ts
+function r(e, t, n, r, i, a) {
+	let o = n.encodeConversation(e, !1, !!a);
+	if (o) {
+		let e = Array.isArray(o) ? o : o.tokens;
+		r.total += e.length;
+		let n = t[t.length - 1], s = a ? a[a.length - 1] : null;
+		if (r.offset + e.length > n.length) {
+			let c = n.length - r.offset;
+			n.set(e.slice(0, c), r.offset);
+			let l = e.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(e.slice(c), 0), t.push(u), a && s && !Array.isArray(o)) {
+				s.set(o.mask.slice(0, c).map((e) => +!!e), r.offset);
+				let e = new Uint8Array(u.length);
+				e.set(o.mask.slice(c).map((e) => +!!e), 0), a.push(e);
+			}
+			r.offset = e.length - c;
+		} else n.set(e, r.offset), s && !Array.isArray(o) && s.set(o.mask.map((e) => +!!e), r.offset), r.offset += e.length;
+	}
+}
+async function i(i, a, o) {
+	await t("training-tokens");
+	let s = await e("training-tokens", a.id, a.datasetID ?? "", o);
+	await t("validation-tokens");
+	let c = o?.validationSplit && o.validationSplit > 0 ? await e("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
+	}, 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 = 0, g = o?.cb, _ = o?.validationSeed === void 0 ? Math.random : n(o.validationSeed);
+	for (; h < i.length;) await i[h++].begin((e) => {
+		let t = o?.validationSplit && o.validationSplit > 0 && _() < o.validationSplit, n = t ? f : l, i = t ? m : d, h = t ? p : u, g = t ? c : s;
+		r(e, n, a, i, g.shardSize, h || void 0), n.length > 1 && (g.appendShard(n[0], h ? h[0] : void 0), n.shift(), h && h.shift());
+	}, g ? () => g(d.total) : void 0);
+	return l.length === 1 && (l[0] = l[0].subarray(0, d.offset), 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), c.appendShard(f[0], p ? p[0].subarray(0, m.offset) : void 0)), await s.finish(), c && await c.finish(), {
+		trainingTokens: s,
+		validationTokens: f ? c : void 0
+	};
+}
+//#endregion
+export { i as tokensFromStreams };

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;