@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/tokeniser/CharTokeniser.js

@@ -0,0 +1,92 @@
+import { yieldIfNeeded as e } from "../utilities/yielder.js";
+import { n as t, t as n } from "../BaseTokeniser-DSg9zcYq.js";
+//#region lib/tokeniser/CharTokeniser.ts
+var r = ["<eos>", "<unk>"], i = class extends n {
+	vocabSize = 0;
+	eosToken = 0;
+	bosToken = 0;
+	unkToken = 0;
+	vocab = [];
+	cache = /* @__PURE__ */ new Map();
+	_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) => {
+				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) => {
+				this.cache.set(e, t);
+			});
+			else throw Error("Vocab cannot be empty");
+			this._trained = !0;
+		} else this.vocabSize = e, this.vocab = Array(this.vocabSize).fill(""), this.addSpecialTokens(), this.eosToken = this.getSpecialTokenIndex("<eos>"), this.bosToken = this.getSpecialTokenIndex("<bos>") ?? this.eosToken, this.unkToken = this.getSpecialTokenIndex(""), this.vocab.forEach((e, t) => {
+			this.cache.set(e, t);
+		}), this.cache.set("", this.unkToken);
+	}
+	addToken(e, t) {
+		if (this.cache.has(e)) return this.cache.get(e);
+		let n;
+		if (t === void 0 ? (n = this.vocab.indexOf("", this.unkToken + 1), n === -1 && (n = this.vocabSize)) : n = t, n >= this.vocabSize) throw Error("Vocab size exceeded");
+		return this.vocab[n] = e, this.cache.set(e, n), n;
+	}
+	get trained() {
+		return this.vocab.length === this.vocabSize && this._trained;
+	}
+	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) r.forEach((e) => {
+			for (let t of e.content) a.add(t);
+		}), o = await e(o, n, 0);
+		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);
+		}
+		let u = c;
+		if (u !== -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;
+		}
+		return this.cache.clear(), this.vocab.forEach((e, t) => {
+			this.cache.set(e, t);
+		}), this.generateID(), this.emit("trainStatus", "trained"), this.vocabSize;
+	}
+	tokenise(e, t) {
+		if (!this.trained) throw Error("Tokeniser not trained");
+		return e.map((e) => t ? e.split("").map((e) => this.cache.get(e) ?? this.unkToken) : e.split("").map((e) => {
+			let t = this.cache.get(e);
+			return t === void 0 ? "" : this.vocab[t];
+		}));
+	}
+	detokenise(e) {
+		return e.map((e) => Array.from(e).map((e) => this.vocab[e] || "").join(""));
+	}
+	encode(e) {
+		return this.tokenise([e], !0)[0];
+	}
+	decode(e) {
+		return this.detokenise([e])[0];
+	}
+	getVocab() {
+		return this.vocab;
+	}
+	getMerges() {
+		return [];
+	}
+	async createTrainingData(e, t = 5) {
+		let n = await this.tokenise(e, !0), r = [], i = [];
+		for (let e = 0; e < n.length - t; e++) r.push(...n[e].slice(0, t)), i.push(n[e + 1][0]);
+		return [r, i];
+	}
+};
+//#endregion
+export { i as default };

package/dist/tokeniser/bpe.d.ts

@@ -0,0 +1,28 @@
+import { default as BaseTokeniser } from './BaseTokeniser';
+import { Conversation } from './type';
+export default class BPETokeniser extends BaseTokeniser {
+    private targetSize;
+    private vocab;
+    private vocabIndex;
+    private merges;
+    private pretokenMap;
+    constructor(vocabSize: number);
+    constructor(vocab: string[], merges?: [string, string][]);
+    addToken(token: string, index?: number): number;
+    destroy(): void;
+    get trained(): boolean;
+    get vocabSize(): number;
+    get eosToken(): number;
+    get bosToken(): number;
+    get unkToken(): number;
+    train(text?: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
+    getVocab(): string[];
+    getMerges(): [string, string][];
+    private tokeniseWord;
+    private tokeniseStrings;
+    tokenise(text: string[], numeric: true): number[][];
+    tokenise(text: string[]): string[][];
+    detokenise(tokens: number[][]): string[];
+    encode(text: string): number[];
+    decode(tokens: number[]): string;
+}

package/dist/tokeniser/bpe.js

@@ -0,0 +1,170 @@
+import { yieldIfNeeded as e } from "../utilities/yielder.js";
+import { n as t, t as n } from "../BaseTokeniser-DSg9zcYq.js";
+import r from "../utilities/tokenParse.js";
+//#region lib/tokeniser/bpe.ts
+function i(e, t) {
+	return `${e}-::-${t}`;
+}
+function a(e) {
+	let t = /* @__PURE__ */ new Map();
+	for (let n = 0; n < e.length; n++) {
+		let r = e[n];
+		for (let e = 0; e < r.length - 1; e++) {
+			let a = i(r[e], r[e + 1]), o = t.get(a) || {
+				a: r[e],
+				b: r[e + 1],
+				count: 0,
+				instances: /* @__PURE__ */ new Set()
+			};
+			o.count += 1, o.instances.add(n), t.set(a, o);
+		}
+	}
+	return {
+		pairs: t,
+		tokens: e
+	};
+}
+function o(e, t, n, r, a) {
+	let o = i(t, n);
+	if (e.pairs.has(o)) {
+		let t = e.pairs.get(o);
+		t.count += a, a > 0 ? t.instances.add(r) : t.count <= 0 ? e.pairs.delete(o) : t.instances.delete(r);
+	} else e.pairs.set(o, {
+		a: t,
+		b: n,
+		count: a,
+		instances: new Set([r])
+	});
+}
+function s(e) {
+	let t = null, n = 0;
+	for (let r of e.pairs.values()) r.count > n && (n = r.count, t = r);
+	return t;
+}
+function c(e, t) {
+	return e.map((e) => {
+		let n = [];
+		for (let r = 0; r < e.length; r++) r < e.length - 1 && e[r] === t[0] && e[r + 1] === t[1] ? (n.push(t[0] + t[1]), r++) : n.push(e[r]);
+		return n;
+	});
+}
+function l(e, t) {
+	t.instances.forEach((n) => {
+		let r = e.tokens[n], i = [];
+		for (let a = 0; a < r.length; a++) if (a < r.length - 1 && r[a] === t.a && r[a + 1] === t.b) {
+			let s = t.a + t.b;
+			i.push(s), a > 0 && (o(e, r[a - 1], t.a, n, -1), o(e, r[a - 1], s, n, 1)), a++, a < r.length - 1 && (o(e, t.b, r[a + 1], n, -1), o(e, s, r[a + 1], n, 1));
+		} else i.push(r[a]);
+		e.tokens[n] = i;
+	}), e.pairs.delete(i(t.a, t.b));
+}
+var u = class extends n {
+	targetSize;
+	vocab = /* @__PURE__ */ new Set();
+	vocabIndex = /* @__PURE__ */ new Map();
+	merges = [];
+	pretokenMap = /* @__PURE__ */ new Map();
+	constructor(e, n) {
+		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) => {
+			let n = e.indexOf(t);
+			n !== -1 && this.addSpecialToken(t, n);
+		})) : (this.addSpecialTokens(), this.targetSize = e);
+	}
+	addToken(e, t) {
+		if (this.vocab.has(e)) return this.vocabIndex.get(e);
+		{
+			this.vocab.add(e);
+			let n = t === void 0 ? this.vocab.size - 1 : t;
+			return this.vocabIndex.set(e, n), n;
+		}
+	}
+	destroy() {
+		this.vocab.clear(), this.vocabIndex.clear(), this.merges = [], this.pretokenMap.clear();
+	}
+	get trained() {
+		return this.vocab.size > t.length && this.vocab.size <= this.targetSize;
+	}
+	get vocabSize() {
+		return this.vocab.size;
+	}
+	get eosToken() {
+		return this.vocabIndex.get("<eos>") ?? 0;
+	}
+	get bosToken() {
+		return this.vocabIndex.get("<bos>") ?? 0;
+	}
+	get unkToken() {
+		return this.vocabIndex.get("") ?? 1;
+	}
+	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);
+		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 (; this.vocab.size < this.targetSize && this.merges.length < this.targetSize;) {
+			let t = s(m);
+			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);
+		}
+		f.forEach((e, t) => {
+			let n = p[t];
+			this.pretokenMap.set(e, n);
+		}), this.vocabIndex.clear();
+		let h = 0;
+		for (let e of this.vocab.keys()) this.vocabIndex.set(e, h++);
+		return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
+	}
+	getVocab() {
+		return Array.from(this.vocab);
+	}
+	getMerges() {
+		return this.merges;
+	}
+	tokeniseWord(e) {
+		let t = Array.from(e);
+		return this.merges.forEach((e) => {
+			t = c([t], e)[0];
+		}), this.pretokenMap.set(e, t), t;
+	}
+	tokeniseStrings(e) {
+		return e.map((e) => r(e).map((e) => this.pretokenMap.has(e) ? this.pretokenMap.get(e) : this.tokeniseWord(e)).flat(1));
+	}
+	tokenise(e, t) {
+		let n = this.tokeniseStrings(e);
+		return t ? n.map((e) => e.map((e) => this.vocabIndex.get(e) ?? this.unkToken)) : n.map((e) => e.map((e) => this.vocab.has(e) ? e : ""));
+	}
+	detokenise(e) {
+		let t = this.getVocab();
+		return e.map((e) => e.map((e) => t[e]).join(""));
+	}
+	encode(e) {
+		return this.tokenise([e], !0)[0];
+	}
+	decode(e) {
+		return this.detokenise([e])[0];
+	}
+};
+//#endregion
+export { u as default };

package/dist/tokeniser/messages.d.ts

@@ -0,0 +1,61 @@
+interface TrainMessage {
+    type: 'train';
+    id: number;
+    text: string[];
+    vocabSize: number;
+}
+interface TrainResponse {
+    type: 'trainResponse';
+    id: number;
+    vocabSize: number;
+}
+interface TrainStatusMessage {
+    type: 'trainStatus';
+    id: number;
+    progress: number;
+    vocabSize: number;
+}
+interface TokeniseMessage {
+    type: 'tokenise';
+    id: number;
+    numeric?: boolean;
+    text: string[];
+}
+interface TokeniseResponse {
+    type: 'tokeniseResponse';
+    id: number;
+    numeric: boolean;
+    tokens: string[][] | number[][];
+}
+interface DetokeniseMessage {
+    type: 'detokenise';
+    id: number;
+    tokens: number[][];
+}
+interface DetokeniseResponse {
+    type: 'detokeniseResponse';
+    id: number;
+    text: string[];
+}
+interface TokensMessage {
+    type: 'tokens';
+    id: number;
+}
+interface TokensResponse {
+    type: 'tokensResponse';
+    id: number;
+    tokens: string[];
+}
+interface BuildTrainingDataMessage {
+    type: 'buildTrainingData';
+    id: number;
+    text: string[];
+    windowSize: number;
+}
+interface BuildTrainingDataResponse {
+    type: 'buildTrainingDataResponse';
+    id: number;
+    trainingData: [number[], number[]];
+}
+export type TokeniserMessage = TrainMessage | TrainResponse | TrainStatusMessage | TokeniseMessage | DetokeniseMessage | TokeniseResponse | DetokeniseResponse | TokensMessage | TokensResponse | BuildTrainingDataMessage | BuildTrainingDataResponse;
+export {};

package/dist/tokeniser/type.d.ts

@@ -0,0 +1,34 @@
+import { default as EE } from 'eventemitter3';
+export type Roles = 'user' | 'assistant' | 'system' | 'text';
+export interface Conversation {
+    role: Roles;
+    content: string;
+}
+export interface ITokeniser extends EE<'trainStatus'> {
+    id: string;
+    datasetID?: string;
+    train(text: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
+    getVocab(): string[];
+    getMerges(): [string, string][];
+    destroy(): void;
+    encode(text: string): number[];
+    encodeConversation(conversation: Conversation[], completion?: boolean): number[];
+    encodeConversation(conversation: Conversation[], completion: boolean, masking: boolean): {
+        tokens: number[];
+        mask: boolean[];
+    };
+    encodeConversation(conversation: Conversation[], completion?: boolean, masking?: boolean): number[] | {
+        tokens: number[];
+        mask: boolean[];
+    };
+    encodeSequence(text: string): number[];
+    encodeAsSequence(conversation: Conversation[], completion?: boolean): number[];
+    decode(tokens: number[] | Uint16Array): string;
+    decodeConversation(tokens: number[] | Uint16Array): Conversation[];
+    vocabSize: number;
+    eosToken: number;
+    bosToken: number;
+    trained: boolean;
+    getSpecialTokenIndex(token: string): number | undefined;
+    isSpecialToken(index: number): boolean;
+}

package/dist/training/AdamW.d.ts

@@ -0,0 +1,36 @@
+import { Optimizer, Tensor } from '@tensorflow/tfjs-core';
+import { ConfigDict, Serializable, SerializableConstructor } from '@tensorflow/tfjs-core/dist/serialization';
+import { NamedTensor, NamedVariableMap } from '@tensorflow/tfjs-core/dist/tensor_types';
+import { default as LRScheduler } from './LRScheduler';
+import { AdamWOptimizerConfig } from './types';
+export declare class AdamWOptimizer extends Optimizer {
+    private config;
+    readonly className = "AdamW";
+    private accBeta1;
+    private accBeta2;
+    private accumulatedMoments;
+    protected learningRate: number;
+    protected beta1: number;
+    protected beta2: number;
+    protected lossScaling: number;
+    protected weightDecay: number;
+    protected epsilon: number | null;
+    protected lrScheduler: LRScheduler;
+    protected clipNorm?: number;
+    protected orthGradEpsilon: number;
+    protected orthGrad: boolean;
+    constructor(config: AdamWOptimizerConfig);
+    get lr(): number;
+    saveMoments(): Promise<ArrayBuffer>;
+    loadMoments(momentData: ArrayBuffer): Promise<void>;
+    serializeConfig(): AdamWOptimizerConfig;
+    private orthogonalizeGradient;
+    updateConfig(newConfig: Partial<AdamWOptimizerConfig>): void;
+    applyGradients(variableGradients: NamedVariableMap | NamedTensor[]): Tensor;
+    dispose(): void;
+    getWeights(): Promise<NamedTensor[]>;
+    setWeights(weightValues: NamedTensor[]): Promise<void>;
+    getConfig(): ConfigDict;
+    /** @nocollapse */
+    static fromConfig<T extends Serializable>(cls: SerializableConstructor<T>, config: ConfigDict): T;
+}

package/dist/training/AdamW.js

@@ -0,0 +1,128 @@
+import { _n as e, c as t, di as n, ii as r, kt as i, ni as a } from "../dist-BewPQWjc.js";
+import { load_safetensors as o, save_safetensors as s } from "../utilities/safetensors.js";
+import { adamAdjust as c } from "../ops/adamAdjust.js";
+import { adamMoments as l } from "../ops/adamMoments.js";
+import u from "./LRScheduler.js";
+import { clipScale as d } from "../ops/globalNorm.js";
+//#region lib/training/AdamW.ts
+var f = class extends t {
+	config;
+	className = "AdamW";
+	accBeta1 = 0;
+	accBeta2 = 0;
+	accumulatedMoments = [];
+	learningRate;
+	beta1;
+	beta2;
+	lossScaling;
+	weightDecay;
+	epsilon = null;
+	lrScheduler;
+	clipNorm;
+	orthGradEpsilon = 1e-30;
+	orthGrad;
+	constructor(e) {
+		super(), this.config = e, this.accBeta1 = e.accBeta1 ?? e.beta1, this.accBeta2 = e.accBeta2 ?? e.beta2, this.learningRate = e.learningRate, this.beta1 = e.beta1, this.beta2 = e.beta2, this.weightDecay = e.weightDecay, this.lossScaling = e.lossScaling, this.clipNorm = e.clipNorm, this.orthGrad = e.orthoGrad ?? !1, e.epsilon === null || e.epsilon === void 0 ? this.epsilon = r().backend.epsilon() : this.epsilon = e.epsilon, this.lrScheduler = new u(e.learningRate, e);
+	}
+	get lr() {
+		return this.learningRate;
+	}
+	saveMoments() {
+		let e = {};
+		return this.accumulatedMoments.forEach((t) => {
+			e[t.originalName] = t.variable;
+		}), s(e);
+	}
+	async loadMoments(e) {
+		let t = await o(e);
+		Object.entries(t).forEach(([e, t]) => {
+			let n = t.variable(!1);
+			this.accumulatedMoments.push({
+				originalName: e,
+				variable: n
+			});
+		});
+	}
+	serializeConfig() {
+		return {
+			learningRate: this.learningRate,
+			beta1: this.beta1,
+			beta2: this.beta2,
+			accBeta1: this.accBeta1,
+			accBeta2: this.accBeta2,
+			epsilon: this.epsilon ?? void 0,
+			weightDecay: this.weightDecay,
+			lossScaling: this.lossScaling,
+			clipNorm: this.clipNorm,
+			orthoGrad: this.orthGrad,
+			...this.lrScheduler.serializeConfig()
+		};
+	}
+	orthogonalizeGradient(e, t) {
+		return n(() => {
+			let n = e.reshape([-1]), r = t.reshape([-1]), i = n.mul(n).sum().add(this.orthGradEpsilon), a = n.mul(r).sum().div(i), o = r.sub(n.mul(a)), s = r.norm(), c = o.norm().add(this.orthGradEpsilon);
+			return o.mul(s.div(c)).reshape(t.shape);
+		});
+	}
+	updateConfig(e) {
+		let t = {
+			...this.config,
+			...e
+		};
+		this.learningRate = t.learningRate, this.beta1 = t.beta1, this.beta2 = t.beta2, this.weightDecay = t.weightDecay, this.lossScaling = t.lossScaling, this.epsilon = t.epsilon ?? this.epsilon, this.clipNorm = t.clipNorm, this.lrScheduler.updateConfig(t, t.learningRate);
+	}
+	applyGradients(t) {
+		let a = this.lrScheduler.getNextLR();
+		this.learningRate = a;
+		let o = Array.isArray(t) ? t.map((e) => e.name) : Object.keys(t), s = n(() => {
+			let a = 1 - this.accBeta1, s = 1 - this.accBeta2, u;
+			return u = this.clipNorm === void 0 ? e(1 / this.lossScaling) : d(o.map((e, n) => Array.isArray(t) ? t[n].tensor : t[e]), 1 / this.lossScaling, this.clipNorm), o.forEach((e, o) => {
+				let d = r().registeredVariables[e];
+				this.accumulatedMoments[o] ?? (this.accumulatedMoments[o] = {
+					originalName: `${e}/m`,
+					variable: n(() => i([...d.shape, 2]).variable(!1))
+				});
+				let f = Array.isArray(t) ? t[o].tensor : t[e];
+				if (f == null) return;
+				let p = this.orthGrad ? this.orthogonalizeGradient(d, f) : f, m = this.accumulatedMoments[o].variable, h = l(m, p, this.beta1, this.beta2, u);
+				m.assign(h), this.orthGrad && p.dispose();
+				let g = c(h, d, a, s, this.epsilon ?? 1e-8, this.learningRate, d.shape.length > 1 ? this.weightDecay : 0);
+				d.assign(g);
+			}), this.accBeta1 *= this.beta1, this.accBeta2 *= this.beta2, u;
+		});
+		return this.incrementIterations(), s;
+	}
+	dispose() {
+		this.accumulatedMoments != null && a(this.accumulatedMoments.map((e) => e.variable));
+	}
+	async getWeights() {
+		let e = [...this.accumulatedMoments];
+		return [await this.saveIterations()].concat(e.map((e) => ({
+			name: e.originalName,
+			tensor: e.variable
+		})));
+	}
+	async setWeights(e) {
+		e = await this.extractIterations(e), n(() => {
+			this.accBeta1 = this.beta1 ** +(this.iterations_ + 1), this.accBeta2 = this.beta2 ** +(this.iterations_ + 1);
+		});
+		let t = e.length / 2;
+		this.accumulatedMoments = e.slice(0, t).map((e) => ({
+			originalName: e.name,
+			variable: e.tensor.variable(!1)
+		}));
+	}
+	getConfig() {
+		return {
+			learningRate: this.learningRate,
+			beta1: this.beta1,
+			beta2: this.beta2,
+			epsilon: this.epsilon
+		};
+	}
+	static fromConfig(e, t) {
+		return new e(t.learningRate, t.beta1, t.beta2, t.epsilon);
+	}
+};
+//#endregion
+export { f as AdamWOptimizer };

package/dist/training/BasicTrainer.d.ts

@@ -0,0 +1,63 @@
+import { ITokeniser } from '../tokeniser/type';
+import { Scalar, Tensor } from '@tensorflow/tfjs-core';
+import { Dataset } from '@tensorflow/tfjs-data';
+import { default as Model, ModelForwardAttributes } from '../../models/model';
+import { AdamWOptimizerConfig, TrainingLogEntry, TrainingMetrics, TrainingOptions, TrainingState } from './types';
+import { AdamWOptimizer } from './AdamW';
+export default class BasicTrainer {
+    tokenizer: ITokeniser;
+    model: Model<ModelForwardAttributes>;
+    optimizer: AdamWOptimizer;
+    protected running: boolean;
+    protected lastState?: TrainingState;
+    protected _gradientCheckpointing: boolean;
+    protected _mixedPrecision: boolean;
+    protected maskedLoss: boolean;
+    protected optimizerConfig: AdamWOptimizerConfig;
+    protected metrics: Set<TrainingMetrics>;
+    protected _labelSmoothing: number;
+    protected _layerDrop: number;
+    protected _dropout: number;
+    constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, optConfig?: Partial<AdamWOptimizerConfig>, optimizer?: AdamWOptimizer);
+    setLossMasking(): void;
+    setGradientCheckpointing(enabled: boolean): void;
+    setMixedPrecision(enabled: boolean): void;
+    setLabelSmoothing(smoothing: number): void;
+    setDropout(dropout: number): void;
+    setLayerDrop(layerDrop: number): void;
+    setLearningRate(learningRate: number): void;
+    setMetrics(metrics: TrainingMetrics[]): void;
+    reset(): void;
+    stop(): void;
+    get isRunning(): boolean;
+    getOptimizer(): AdamWOptimizer;
+    updateOptimizer(config?: Partial<AdamWOptimizerConfig>): void;
+    resumeFromLog(log: TrainingLogEntry): void;
+    protected trainStep(state: Partial<TrainingState>, batch: {
+        xs: Tensor;
+        ys: Tensor;
+    }, dummy?: boolean, keepGrads?: boolean): Scalar;
+    private dummyPass;
+    dispose(): void;
+    private createEmptyState;
+    stepDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>): Promise<{
+        log: TrainingLogEntry;
+    }>;
+    private performLogging;
+    trainOnDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>): Promise<{
+        losses: number[];
+        validationLosses: number[];
+    }>;
+}