@genai-fi/nanogpt 0.19.1 → 0.20.1

package/dist/tokeniser/bpe.js

@@ -1,173 +1,170 @@
-import { yieldIfNeeded as p } from "../utilities/yielder.js";
-import m from "../utilities/tokenParse.js";
-import T, { SPECIALS as S } from "./BaseTokeniser.js";
-function g(o, e) {
-  return `${o}-::-${e}`;
+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 y(o) {
-  const e = /* @__PURE__ */ new Map();
-  for (let s = 0; s < o.length; s++) {
-    const t = o[s];
-    for (let n = 0; n < t.length - 1; n++) {
-      const r = g(t[n], t[n + 1]), a = e.get(r) || {
-        a: t[n],
-        b: t[n + 1],
-        count: 0,
-        instances: /* @__PURE__ */ new Set()
-      };
-      a.count += 1, a.instances.add(s), e.set(r, a);
-    }
-  }
-  return { pairs: e, tokens: o };
+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 f(o, e, s, t, n) {
-  const r = g(e, s);
-  if (o.pairs.has(r)) {
-    const a = o.pairs.get(r);
-    a.count += n, n > 0 ? a.instances.add(t) : a.count <= 0 ? o.pairs.delete(r) : a.instances.delete(t);
-  } else
-    o.pairs.set(r, { a: e, b: s, count: n, instances: /* @__PURE__ */ new Set([t]) });
+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 I(o) {
-  let e = null, s = 0;
-  for (const t of o.pairs.values())
-    t.count > s && (s = t.count, e = t);
-  return e;
+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 x(o, e) {
-  return o.map((s) => {
-    const t = [];
-    for (let n = 0; n < s.length; n++)
-      n < s.length - 1 && s[n] === e[0] && s[n + 1] === e[1] ? (t.push(e[0] + e[1]), n++) : t.push(s[n]);
-    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 A(o, e) {
-  e.instances.forEach((s) => {
-    const t = o.tokens[s], n = [];
-    for (let r = 0; r < t.length; r++)
-      if (r < t.length - 1 && t[r] === e.a && t[r + 1] === e.b) {
-        const a = e.a + e.b;
-        n.push(a), r > 0 && (f(o, t[r - 1], e.a, s, -1), f(o, t[r - 1], a, s, 1)), r++, r < t.length - 1 && (f(o, e.b, t[r + 1], s, -1), f(o, a, t[r + 1], s, 1));
-      } else
-        n.push(t[r]);
-    o.tokens[s] = n;
-  }), o.pairs.delete(g(e.a, e.b));
+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));
 }
-class P extends T {
-  targetSize;
-  vocab = /* @__PURE__ */ new Set();
-  vocabIndex = /* @__PURE__ */ new Map();
-  merges = [];
-  pretokenMap = /* @__PURE__ */ new Map();
-  constructor(e, s) {
-    super(), Array.isArray(e) ? (e.forEach((t, n) => {
-      this.vocab.add(t), this.vocabIndex.set(t, n);
-    }), s && (this.merges = s), this.targetSize = e.length, S.forEach((t) => {
-      const n = e.indexOf(t);
-      n !== -1 && this.addSpecialToken(t, n);
-    })) : (this.addSpecialTokens(), this.targetSize = e);
-  }
-  addToken(e, s) {
-    if (this.vocab.has(e))
-      return this.vocabIndex.get(e);
-    {
-      this.vocab.add(e);
-      const t = s !== void 0 ? s : this.vocab.size - 1;
-      return this.vocabIndex.set(e, t), t;
-    }
-  }
-  destroy() {
-    this.vocab.clear(), this.vocabIndex.clear(), this.merges = [], this.pretokenMap.clear();
-  }
-  get trained() {
-    return this.vocab.size > S.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(e = [], s, t) {
-    this.datasetID = t;
-    let n = performance.now();
-    const r = new Array(e.length);
-    for (let i = 0; i < e.length; i++) {
-      const h = e[i], l = new Array(h.length);
-      for (let d = 0; d < h.length; d++)
-        l[d] = m(h[d].content);
-      n = await p(n, s, this.vocab.size), r[i] = l;
-    }
-    const a = r.flat(2), z = new Set(a);
-    this.vocab = /* @__PURE__ */ new Set(), this.pretokenMap.clear(), this.merges = [], this.addSpecialTokens();
-    const b = Array.from(z), v = b.map((i) => Array.from(i).map((l) => (this.vocab.add(l), l))), k = y(v);
-    if (n = await p(n, s, 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.");
-      const i = /* @__PURE__ */ new Map();
-      a.forEach((c) => {
-        Array.from(c).forEach((u) => {
-          i.set(u, (i.get(u) || 0) + 1);
-        });
-      });
-      const h = Array.from(i.entries()).sort((c, u) => u[1] - c[1]);
-      this.vocab = /* @__PURE__ */ new Set(), this.addSpecialTokens(), h.slice(0, this.targetSize - this.vocab.size).map(([c]) => c).forEach((c) => this.vocab.add(c)), this.vocabIndex.clear();
-      let d = 0;
-      for (const c of this.vocab.keys())
-        this.vocabIndex.set(c, d++);
-      return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
-    }
-    for (; this.vocab.size < this.targetSize && this.merges.length < this.targetSize; ) {
-      const i = I(k);
-      if (!i)
-        break;
-      this.merges.push([i.a, i.b]), this.vocab.add(i.a + i.b), A(k, i), n = await p(n, s, this.vocab.size);
-    }
-    b.forEach((i, h) => {
-      const l = v[h];
-      this.pretokenMap.set(i, l);
-    }), this.vocabIndex.clear();
-    let w = 0;
-    for (const i of this.vocab.keys())
-      this.vocabIndex.set(i, w++);
-    return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
-  }
-  getVocab() {
-    return Array.from(this.vocab);
-  }
-  getMerges() {
-    return this.merges;
-  }
-  tokeniseWord(e) {
-    let s = Array.from(e);
-    return this.merges.forEach((t) => {
-      s = x([s], t)[0];
-    }), this.pretokenMap.set(e, s), s;
-  }
-  tokeniseStrings(e) {
-    return e.map((s) => m(s).map((r) => this.pretokenMap.has(r) ? this.pretokenMap.get(r) : this.tokeniseWord(r)).flat(1));
-  }
-  tokenise(e, s) {
-    const t = this.tokeniseStrings(e);
-    return s ? t.map((n) => n.map((r) => this.vocabIndex.get(r) ?? this.unkToken)) : t.map((n) => n.map((r) => this.vocab.has(r) ? r : ""));
-  }
-  detokenise(e) {
-    const s = this.getVocab();
-    return e.map((n) => n.map((r) => s[r]).join(""));
-  }
-  encode(e) {
-    return this.tokenise([e], !0)[0];
-  }
-  decode(e) {
-    return this.detokenise([e])[0];
-  }
-}
-export {
-  P as default
+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.js

@@ -1 +0,0 @@
-

package/dist/tokeniser/type.js

@@ -1 +0,0 @@
-

package/dist/training/AdamW.js

@@ -1,138 +1,128 @@
-import { adamAdjust as B } from "../ops/adamAdjust.js";
-import { adamMoments as N } from "../ops/adamMoments.js";
-import { O as S, h as b, t as c, a as M, d as w } from "../index-CUXkjxiT.js";
-import R from "./LRScheduler.js";
-import { clipScale as f } from "../ops/globalNorm.js";
-import { save_safetensors as v, load_safetensors as A } from "../utilities/safetensors.js";
-import { z as O } from "../zeros-DvZpK8s6.js";
-class _ extends S {
-  constructor(t) {
-    super(), this.config = t, this.accBeta1 = t.accBeta1 ?? t.beta1, this.accBeta2 = t.accBeta2 ?? t.beta2, this.learningRate = t.learningRate, this.beta1 = t.beta1, this.beta2 = t.beta2, this.weightDecay = t.weightDecay, this.lossScaling = t.lossScaling, this.clipNorm = t.clipNorm, this.orthGrad = t.orthoGrad ?? !1, t.epsilon === null || t.epsilon === void 0 ? this.epsilon = b().backend.epsilon() : this.epsilon = t.epsilon, this.lrScheduler = new R(t.learningRate, t);
-  }
-  className = "AdamW";
-  accBeta1 = 0;
-  accBeta2 = 0;
-  accumulatedMoments = [];
-  learningRate;
-  beta1;
-  beta2;
-  lossScaling;
-  weightDecay;
-  epsilon = null;
-  lrScheduler;
-  clipNorm;
-  orthGradEpsilon = 1e-30;
-  orthGrad;
-  get lr() {
-    return this.learningRate;
-  }
-  saveMoments() {
-    const t = {};
-    return this.accumulatedMoments.forEach((e) => {
-      t[e.originalName] = e.variable;
-    }), v(t);
-  }
-  async loadMoments(t) {
-    const e = await A(t);
-    Object.entries(e).forEach(([a, s]) => {
-      const n = s.variable(!1);
-      this.accumulatedMoments.push({ originalName: a, 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(t, e) {
-    return c(() => {
-      const a = t.reshape([-1]), s = e.reshape([-1]), n = a.mul(a).sum().add(this.orthGradEpsilon), h = a.mul(s).sum().div(n), o = s.sub(a.mul(h)), l = s.norm(), i = o.norm().add(this.orthGradEpsilon);
-      return o.mul(l.div(i)).reshape(e.shape);
-    });
-  }
-  updateConfig(t) {
-    const e = { ...this.config, ...t };
-    this.learningRate = e.learningRate, this.beta1 = e.beta1, this.beta2 = e.beta2, this.weightDecay = e.weightDecay, this.lossScaling = e.lossScaling, this.epsilon = e.epsilon ?? this.epsilon, this.clipNorm = e.clipNorm, this.lrScheduler.updateConfig(e, e.learningRate);
-  }
-  applyGradients(t) {
-    const e = this.lrScheduler.getNextLR();
-    this.learningRate = e;
-    const a = Array.isArray(t) ? t.map((n) => n.name) : Object.keys(t), s = c(() => {
-      const n = 1 - this.accBeta1, h = 1 - this.accBeta2;
-      let o;
-      if (this.clipNorm !== void 0) {
-        const l = a.map((i, r) => Array.isArray(t) ? t[r].tensor : t[i]);
-        o = f(l, 1 / this.lossScaling, this.clipNorm);
-      } else
-        o = M(1 / this.lossScaling);
-      return a.forEach((l, i) => {
-        const r = b().registeredVariables[l], p = !1;
-        this.accumulatedMoments[i] == null && (this.accumulatedMoments[i] = {
-          originalName: `${l}/m`,
-          variable: c(() => O([...r.shape, 2]).variable(p))
-        });
-        const m = Array.isArray(t) ? t[i].tensor : t[l];
-        if (m == null)
-          return;
-        const u = this.orthGrad ? this.orthogonalizeGradient(r, m) : m, d = this.accumulatedMoments[i].variable, g = N(d, u, this.beta1, this.beta2, o);
-        d.assign(g), this.orthGrad && u.dispose();
-        const y = B(
-          g,
-          r,
-          n,
-          h,
-          this.epsilon ?? 1e-8,
-          this.learningRate,
-          // Only apply weight decay if the variable is multi-dimensional (e.g. weights, not biases)
-          r.shape.length > 1 ? this.weightDecay : 0
-        );
-        r.assign(y);
-      }), this.accBeta1 = this.accBeta1 * this.beta1, this.accBeta2 = this.accBeta2 * this.beta2, o;
-    });
-    return this.incrementIterations(), s;
-  }
-  dispose() {
-    this.accumulatedMoments != null && w(this.accumulatedMoments.map((t) => t.variable));
-  }
-  async getWeights() {
-    const t = [...this.accumulatedMoments];
-    return [await this.saveIterations()].concat(
-      t.map((e) => ({ name: e.originalName, tensor: e.variable }))
-    );
-  }
-  async setWeights(t) {
-    t = await this.extractIterations(t), c(() => {
-      this.accBeta1 = Math.pow(this.beta1, this.iterations_ + 1), this.accBeta2 = Math.pow(this.beta2, this.iterations_ + 1);
-    });
-    const e = t.length / 2, a = !1;
-    this.accumulatedMoments = t.slice(0, e).map((s) => ({
-      originalName: s.name,
-      variable: s.tensor.variable(a)
-    }));
-  }
-  getConfig() {
-    return {
-      learningRate: this.learningRate,
-      beta1: this.beta1,
-      beta2: this.beta2,
-      epsilon: this.epsilon
-    };
-  }
-  /** @nocollapse */
-  static fromConfig(t, e) {
-    return new t(e.learningRate, e.beta1, e.beta2, e.epsilon);
-  }
-}
-export {
-  _ as AdamWOptimizer
+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

@@ -1,7 +1,7 @@
 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 { default as Model, ModelForwardAttributes } from '../../models/model';
 import { AdamWOptimizerConfig, TrainingLogEntry, TrainingMetrics, TrainingOptions, TrainingState } from './types';
 import { AdamWOptimizer } from './AdamW';
 export default class BasicTrainer {