@genai-fi/nanogpt 0.21.0 → 0.23.0

package/dist/tokeniser/BaseTokeniser.d.ts

@@ -1,3 +1,4 @@
+import { ConversationStream } from '../../data/stream';
 import { Conversation, ITokeniser } from './type';
 import { default as EE } from 'eventemitter3';
 export declare const SPECIALS: string[];
@@ -15,7 +16,7 @@ export default abstract class BaseTokeniser extends EE<'trainStatus'> implements
     protected addSpecialTokens(): void;
     protected addSpecialToken(token: string, index: number): void;
     generateID(): void;
-    abstract train(text: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
+    abstract train(text: ConversationStream[], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
     abstract getVocab(): string[];
     abstract getMerges(): [string, string][];
     abstract destroy(): void;

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

@@ -1,5 +1,5 @@
 import { default as BaseTokeniser } from './BaseTokeniser';
-import { Conversation } from './type';
+import { ConversationStream } from '../../data/stream';
 export default class CharTokeniser extends BaseTokeniser {
     vocabSize: number;
     eosToken: number;
@@ -12,7 +12,7 @@ export default class CharTokeniser extends BaseTokeniser {
     addToken(token: string, index?: number): number;
     get trained(): boolean;
     destroy(): void;
-    train(text: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
+    train(text: ConversationStream[], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
     tokenise(text: string[], numeric: true): number[][];
     tokenise(text: string[]): string[][];
     detokenise(tokens: (number[] | Uint16Array)[]): string[];

package/dist/tokeniser/CharTokeniser.js

@@ -1,7 +1,7 @@
 import { yieldIfNeeded as e } from "../utilities/yielder.js";
-import { n as t, t as n } from "../BaseTokeniser-C9TSv4th.js";
+import t, { SPECIALS as n } from "./BaseTokeniser.js";
 //#region lib/tokeniser/CharTokeniser.ts
-var r = ["<eos>", "<unk>"], i = class extends n {
+var r = ["<eos>", "<unk>"], i = class extends t {
 	vocabSize = 0;
 	eosToken = 0;
 	bosToken = 0;
@@ -11,7 +11,7 @@ var r = ["<eos>", "<unk>"], i = class extends n {
 	_trained = !1;
 	constructor(e) {
 		if (super(), Array.isArray(e)) {
-			if (this.vocab = e, this.vocab.length > 0) this.vocabSize = this.vocab.length, t.forEach((e) => {
+			if (this.vocab = e, this.vocab.length > 0) this.vocabSize = this.vocab.length, n.forEach((e) => {
 				let t = this.vocab.indexOf(e);
 				t !== -1 && this.addSpecialToken(e, t);
 			}), this.eosToken = this.getSpecialTokenIndex("<eos>"), this.bosToken = this.getSpecialTokenIndex("<bos>") ?? this.eosToken, this.unkToken = this.getSpecialTokenIndex("") ?? -1, this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("<unk>")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("<pad>")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("_")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf(" ")), this.unkToken === -1 && (this.unkToken = this.eosToken), this.vocab = this.vocab.map((e) => e === "<pad>" ? "" : e), this.vocab.forEach((e, t) => {
@@ -38,19 +38,15 @@ var r = ["<eos>", "<unk>"], i = class extends n {
 	async train(t, n, i) {
 		this.datasetID = i;
 		let a = /* @__PURE__ */ new Set(), o = performance.now();
-		for (let r of t) r.forEach((e) => {
-			for (let t of e.content) a.add(t);
-		}), o = await e(o, n, 0);
+		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) {
-			let e = /* @__PURE__ */ new Map();
-			t.forEach((t) => {
-				t.forEach((t) => {
-					for (let n of t.content) e.set(n, (e.get(n) || 0) + 1);
-				});
-			}), s.sort((t, n) => (e.get(t) || 0) - (e.get(n) || 0)), s.splice(0, s.length - l);
-		}
+		if (this._trained = !0, s.length > l) throw Error("too_small_vocab");
 		let u = c;
 		if (u !== -1) {
 			let e = new Set(this.vocab);

package/dist/tokeniser/bpe.d.ts

@@ -1,5 +1,5 @@
 import { default as BaseTokeniser } from './BaseTokeniser';
-import { Conversation } from './type';
+import { ConversationStream } from '../../data/stream';
 export default class BPETokeniser extends BaseTokeniser {
     private targetSize;
     private vocab;
@@ -15,7 +15,7 @@ export default class BPETokeniser extends BaseTokeniser {
     get eosToken(): number;
     get bosToken(): number;
     get unkToken(): number;
-    train(text?: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
+    train(text?: ConversationStream[], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
     getVocab(): string[];
     getMerges(): [string, string][];
     private tokeniseWord;

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 t, { SPECIALS as n } from "./BaseTokeniser.js";
 import r from "../utilities/tokenParse.js";
 //#region lib/tokeniser/bpe.ts
 function i(e, t) {
@@ -58,16 +58,16 @@ 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 t {
 	targetSize;
 	vocab = /* @__PURE__ */ new Set();
 	vocabIndex = /* @__PURE__ */ new Map();
 	merges = [];
 	pretokenMap = /* @__PURE__ */ new Map();
-	constructor(e, n) {
+	constructor(e, t) {
 		super(), Array.isArray(e) ? (e.forEach((e, t) => {
 			this.vocab.add(e), this.vocabIndex.set(e, t);
-		}), n && (this.merges = n), this.targetSize = e.length, t.forEach((t) => {
+		}), t && (this.merges = t), this.targetSize = e.length, n.forEach((t) => {
 			let n = e.indexOf(t);
 			n !== -1 && this.addSpecialToken(t, n);
 		})) : (this.addSpecialTokens(), this.targetSize = e);
@@ -84,7 +84,7 @@ var u = class extends n {
 		this.vocab.clear(), this.vocabIndex.clear(), this.merges = [], this.pretokenMap.clear();
 	}
 	get trained() {
-		return this.vocab.size > t.length && this.vocab.size <= this.targetSize;
+		return this.vocab.size > n.length && this.vocab.size <= this.targetSize;
 	}
 	get vocabSize() {
 		return this.vocab.size;
@@ -100,40 +100,31 @@ var u = class extends n {
 	}
 	async train(t = [], n, i) {
 		this.datasetID = i;
-		let o = performance.now(), c = Array(t.length);
-		for (let i = 0; i < t.length; i++) {
-			let a = t[i], s = Array(a.length);
-			for (let e = 0; e < a.length; e++) s[e] = r(a[e].content);
-			o = await e(o, n, this.vocab.size), c[i] = s;
-		}
-		let u = c.flat(2), d = new Set(u);
+		let o = performance.now(), c = /* @__PURE__ */ new Set();
 		this.vocab = /* @__PURE__ */ new Set(), this.pretokenMap.clear(), this.merges = [], this.addSpecialTokens();
-		let f = Array.from(d), p = f.map((e) => Array.from(e).map((e) => (this.vocab.add(e), e))), m = a(p);
-		if (o = await e(o, n, this.vocab.size), this.vocab.size >= this.targetSize) {
-			console.warn("Initial vocab size is greater than or equal to target size. No merges will be performed.");
-			let e = /* @__PURE__ */ new Map();
-			u.forEach((t) => {
-				Array.from(t).forEach((t) => {
-					e.set(t, (e.get(t) || 0) + 1);
-				});
-			});
-			let t = Array.from(e.entries()).sort((e, t) => t[1] - e[1]);
-			this.vocab = /* @__PURE__ */ new Set(), this.addSpecialTokens(), t.slice(0, this.targetSize - this.vocab.size).map(([e]) => e).forEach((e) => this.vocab.add(e)), this.vocabIndex.clear();
-			let n = 0;
-			for (let e of this.vocab.keys()) this.vocabIndex.set(e, n++);
-			return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
+		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.has(e) || (c.add(e), Array.from(e).forEach((e) => this.vocab.add(e)));
+				}
+				o = await e(o, n, this.vocab.size), a = await t.next();
+			}
 		}
+		let u = Array.from(c), d = u.map((e) => Array.from(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 (; this.vocab.size < this.targetSize && this.merges.length < this.targetSize;) {
-			let t = s(m);
+			let t = s(f);
 			if (!t) break;
-			this.merges.push([t.a, t.b]), this.vocab.add(t.a + t.b), l(m, t), o = await e(o, n, this.vocab.size);
+			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);
 		}
-		f.forEach((e, t) => {
-			let n = p[t];
+		u.forEach((e, t) => {
+			let n = d[t];
 			this.pretokenMap.set(e, n);
 		}), this.vocabIndex.clear();
-		let h = 0;
-		for (let e of this.vocab.keys()) this.vocabIndex.set(e, h++);
+		let p = 0;
+		for (let e of this.vocab.keys()) this.vocabIndex.set(e, p++);
 		return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
 	}
 	getVocab() {

package/dist/tokeniser/type.d.ts

@@ -1,13 +1,14 @@
 import { default as EE } from 'eventemitter3';
+import { ConversationStream } from '../data/stream';
 export type Roles = 'user' | 'assistant' | 'system' | 'text';
 export interface Conversation {
     role: Roles;
     content: string;
 }
-export interface ITokeniser extends EE<'trainStatus'> {
+export interface ITokeniser extends EE<'trainStatus' | 'error'> {
     id: string;
     datasetID?: string;
-    train(text: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
+    train(text: ConversationStream[], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
     getVocab(): string[];
     getMerges(): [string, string][];
     destroy(): void;

package/dist/training/DatasetBuilder.d.ts

@@ -1,6 +1,7 @@
 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;
@@ -8,14 +9,28 @@ export declare function flattenTokensWithMask(textData: Conversation[][], tokeni
 };
 export declare function shuffle(array: Uint32Array): Uint32Array;
 export interface DatasetState {
+    shuffledShards: Uint32Array;
     shuffledIndexes: Uint32Array;
+    lastShardIndexes: Uint32Array;
+    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-B9UGNQT5.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-C0iJT29K.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-B9UGNQT5.js";
+import { t as e } from "../DatasetBuilder-C0iJT29K.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-B9UGNQT5.js";
+import { t as e } from "../DatasetBuilder-C0iJT29K.js";
 import t from "./BasicTrainer.js";
 //#region lib/training/SFTTrainer.ts
 var n = {

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

@@ -1,18 +1,17 @@
-import { Conversation, ITokeniser } from '../../../main';
+import { Conversation, ConversationStream, ITokeniser } from '../../../main';
 import { Task } from './Task';
 export default class ConversationTask extends Task {
-    private rawConvo;
-    private shuffledIndices;
-    private index;
+    private streams;
+    private streamIndex;
+    private currentCursor;
     get length(): number;
-    constructor(conversations: Conversation[][]);
+    constructor(conversations: ConversationStream[]);
     hasMoreConversations(): boolean;
-    nextConversation(): Conversation[] | null;
-    nextTokens(tokeniser: ITokeniser): number[] | null;
-    nextTokens(tokeniser: ITokeniser, masking: boolean): {
+    nextConversation(): Promise<Conversation[] | null>;
+    nextTokens(tokeniser: ITokeniser): Promise<number[] | null>;
+    nextTokens(tokeniser: ITokeniser, masking: boolean): Promise<{
         tokens: number[];
         mask: boolean[];
-    } | null;
-    shuffle(): void;
+    } | null>;
     estimateTokens(tokeniser: ITokeniser): Promise<number>;
 }

package/dist/training/tasks/ConversationTask.js

@@ -1,38 +1,29 @@
-import { i as e } from "../../DatasetBuilder-B9UGNQT5.js";
-import { Task as t } from "./Task.js";
+import { Task as e } from "./Task.js";
 //#region lib/training/tasks/ConversationTask.ts
-var n = class extends t {
-	rawConvo;
-	shuffledIndices = null;
-	index = 0;
+var t = class extends e {
+	streams;
+	streamIndex = 0;
+	currentCursor = null;
 	get length() {
-		return this.rawConvo.length;
+		return this.streams.length;
 	}
 	constructor(e) {
-		super(), this.rawConvo = e;
+		super(), this.streams = e;
 	}
 	hasMoreConversations() {
-		return this.index < this.rawConvo.length;
+		return this.streamIndex < this.streams.length;
 	}
-	nextConversation() {
-		if (this.index >= this.rawConvo.length) return null;
-		let e = this.rawConvo[this.shuffledIndices ? this.shuffledIndices[this.index] : this.index];
-		return this.index++, e;
+	async nextConversation() {
+		return this.streamIndex < this.streams.length ? (this.currentCursor ||= this.streams[this.streamIndex].cursor(), await this.currentCursor.next() || (this.streamIndex++, this.currentCursor = null, this.nextConversation())) : null;
 	}
-	nextTokens(e, t) {
-		let n = this.nextConversation();
+	async nextTokens(e, t) {
+		let n = await this.nextConversation();
 		return n ? e.encodeConversation(n, !1, t) : null;
 	}
-	shuffle() {
-		if (!this.shuffledIndices) {
-			this.shuffledIndices = new Uint32Array(this.rawConvo.length);
-			for (let e = 0; e < this.rawConvo.length; e++) this.shuffledIndices[e] = e;
-		}
-		e(this.shuffledIndices), this.index = 0;
-	}
 	async estimateTokens(e) {
-		return e.encodeConversation(this.rawConvo[0]).length * this.length;
+		let t = await this.streams[0].cursor().next();
+		return t ? e.encodeConversation(t).length * this.length : 0;
 	}
 };
 //#endregion
-export { n as default };
+export { t as default };

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

@@ -6,12 +6,12 @@ export default class PretrainingTask extends Task {
     get length(): number;
     constructor(texts: string[]);
     hasMoreConversations(): boolean;
-    nextConversation(): Conversation[] | null;
-    nextTokens(tokeniser: ITokeniser): number[] | null;
-    nextTokens(tokeniser: ITokeniser, masking: boolean): {
+    nextConversation(): Promise<Conversation[] | null>;
+    nextTokens(tokeniser: ITokeniser): Promise<number[] | null>;
+    nextTokens(tokeniser: ITokeniser, masking: boolean): Promise<{
         tokens: number[];
         mask: boolean[];
-    } | null;
+    } | null>;
     shuffle(): void;
     estimateTokens(tokeniser: ITokeniser): Promise<number>;
 }

package/dist/training/tasks/PretrainingTask.js

@@ -12,7 +12,7 @@ var t = class extends e {
 	hasMoreConversations() {
 		return this.index < this.rawText.length;
 	}
-	nextConversation() {
+	async nextConversation() {
 		if (this.index >= this.rawText.length) return null;
 		let e = {
 			role: "assistant",
@@ -20,7 +20,7 @@ var t = class extends e {
 		};
 		return this.index++, [e];
 	}
-	nextTokens(e, t) {
+	async nextTokens(e, t) {
 		if (this.index >= this.rawText.length) return null;
 		let n = e.encodeSequence(this.rawText[this.index]);
 		return this.index++, t ? {

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

@@ -6,12 +6,12 @@ export default class StartSentenceTask extends Task {
     get length(): number;
     constructor(texts: string[]);
     hasMoreConversations(): boolean;
-    nextConversation(): Conversation[] | null;
-    nextTokens(tokeniser: ITokeniser): number[] | null;
-    nextTokens(tokeniser: ITokeniser, masking: boolean): {
+    nextConversation(): Promise<Conversation[] | null>;
+    nextTokens(tokeniser: ITokeniser): Promise<number[] | null>;
+    nextTokens(tokeniser: ITokeniser, masking: boolean): Promise<{
         tokens: number[];
         mask: boolean[];
-    } | null;
+    } | null>;
     shuffle(): void;
     private conversationFromString;
     estimateTokens(tokeniser: ITokeniser): Promise<number>;

package/dist/training/tasks/StartSentenceTask.js

@@ -12,13 +12,13 @@ var t = class extends e {
 	hasMoreConversations() {
 		return this.index < this.rawText.length;
 	}
-	nextConversation() {
+	async nextConversation() {
 		if (this.index >= this.rawText.length) return null;
 		let e = this.rawText[this.index];
 		return this.index++, this.conversationFromString(e);
 	}
-	nextTokens(e, t) {
-		let n = this.nextConversation();
+	async nextTokens(e, t) {
+		let n = await this.nextConversation();
 		return n ? e.encodeConversation(n, !1, t) : null;
 	}
 	shuffle() {

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

@@ -1,22 +1,29 @@
 import { Conversation, ITokeniser } from '../../../main';
+import { TokenStore } from './TokenStore';
 export declare abstract class Task {
     abstract get length(): number;
     abstract hasMoreConversations(): boolean;
-    abstract nextConversation(): Conversation[] | null;
-    abstract nextTokens(tokeniser: ITokeniser): number[] | null;
-    abstract nextTokens(tokeniser: ITokeniser, masking: boolean): {
+    abstract nextConversation(): Promise<Conversation[] | null>;
+    abstract nextTokens(tokeniser: ITokeniser): Promise<number[] | null>;
+    abstract nextTokens(tokeniser: ITokeniser, masking: boolean): Promise<{
         tokens: number[];
         mask: boolean[];
-    } | null;
-    abstract nextTokens(tokeniser: ITokeniser, masking?: boolean): number[] | {
+    } | null>;
+    abstract nextTokens(tokeniser: ITokeniser, masking?: boolean): Promise<number[] | {
         tokens: number[];
         mask: boolean[];
-    } | null;
-    abstract estimateTokens(tokeniser: ITokeniser): Promise<number>;
-    abstract shuffle(): void;
+    } | null>;
 }
-export declare function tokensFromTasks(tasks: Task[], tokenizer: ITokeniser, cb?: (tokens: number) => void): Promise<Uint16Array>;
-export declare function tokensFromTasks(tasks: Task[], tokenizer: ITokeniser, cb?: (tokens: number) => void, masking?: boolean): Promise<{
-    tokens: Uint16Array;
-    mask: Uint8Array;
+interface TokensFromTasksOptions {
+    masking?: boolean;
+    maxCachedShards?: number;
+    noOPFS?: boolean;
+    shardSize?: number;
+    validationSplit?: number;
+    cb?: (tokens: number) => void;
+}
+export declare function tokensFromTasks(tasks: Task[], tokenizer: ITokeniser, options?: TokensFromTasksOptions): Promise<{
+    trainingTokens: TokenStore;
+    validationTokens?: TokenStore;
 }>;
+export {};