@genai-fi/nanogpt 0.5.2 → 0.5.3

package/dist/training/DatasetBuilder.d.ts

@@ -1,11 +1,14 @@
 import { Tensor } from '@tensorflow/tfjs-core';
 import { ITokeniser } from '../tokeniser/type';
 import { Dataset } from '@tensorflow/tfjs-data';
+export declare const PAGE_FACTOR = 8;
+export declare function flattenTokens(textData: string[], tokenizer: ITokeniser): Promise<number[]>;
 export declare class DatasetBuilder {
     tokenizer: ITokeniser;
     blockSize: number;
+    private pageSize;
     constructor(tokenizer: ITokeniser, blockSize?: number);
-    createTextDataset(textData: string[], batchSize?: number, start?: number, end?: number): Promise<Dataset<{
+    createTextDataset(flatTokens: number[], batchSize?: number, masked?: Set<number>, invertMask?: boolean): Promise<Dataset<{
         xs: Tensor;
         ys: Tensor;
     }>>;

package/dist/training/DatasetBuilder.js

@@ -1,5 +1,5 @@
-import { t as f } from "../index-CnHyhpKc.js";
-import { d as u, i as z } from "../dataset-ZHEPJmED.js";
+import { t as u } from "../index-CnHyhpKc.js";
+import { d as z, i as f } from "../dataset-ZHEPJmED.js";
 import "../index-Tf7vU29b.js";
 /**
  * @license
@@ -18,39 +18,64 @@ import "../index-Tf7vU29b.js";
  *
  * =============================================================================
  */
-function d(a) {
-  return u(async () => {
-    const e = await a();
-    return z(() => e.next());
+function S(c) {
+  return z(async () => {
+    const t = await c();
+    return f(() => t.next());
   });
 }
-class y {
+const p = 8;
+async function y(c, t) {
+  const s = await Promise.all(c.map((n) => t.encode(n))), i = t.eosToken >= 0;
+  return s.map((n) => i ? [...n, t.eosToken] : n).flat();
+}
+class w {
   tokenizer;
   blockSize;
-  constructor(e, s = 128) {
-    this.tokenizer = e, this.blockSize = s;
+  pageSize;
+  constructor(t, s = 128) {
+    this.tokenizer = t, this.blockSize = s, this.pageSize = s * p;
   }
   // Create dataset from text files
-  async createTextDataset(e, s = 32, c = 0, r = 1) {
-    const l = await Promise.all(e.map((t) => this.tokenizer.encode(t))), h = this.tokenizer.eosToken >= 0, n = l.map((t) => h ? [...t, this.tokenizer.eosToken] : t).flat(), i = n.slice(
-      Math.floor(c * n.length),
-      r === 1 ? void 0 : Math.floor(r * n.length)
-    ), k = (function* () {
-      for (; ; ) {
-        const t = Math.floor(Math.random() * (i.length - this.blockSize - 1)), o = i.slice(t, t + this.blockSize), m = i.slice(t + 1, t + this.blockSize + 1);
-        yield { xs: o, ys: m };
-      }
+  async createTextDataset(t, s = 32, i, r) {
+    if (t.length < this.blockSize + 1)
+      throw new Error(`Not enough tokens (${t.length}) for block size ${this.blockSize}`);
+    if (i && i.size > t.length / this.pageSize / 2)
+      throw new Error("Too many masked pages - would leave insufficient training data");
+    const n = (function* () {
+      if (i && r) {
+        const e = Array.from(i);
+        for (; ; ) {
+          const a = Math.floor(Math.random() * e.length), l = Math.floor(Math.random() * this.pageSize), o = e[a] * this.pageSize + l;
+          if (o + this.blockSize + 1 > t.length)
+            continue;
+          const h = t.slice(o, o + this.blockSize), g = t.slice(o + 1, o + this.blockSize + 1);
+          yield { xs: h, ys: g };
+        }
+      } else
+        for (; ; ) {
+          const e = Math.floor(Math.random() * (t.length - this.blockSize - 1));
+          if (i) {
+            const o = Math.floor(e / this.pageSize), h = i.has(o);
+            if (h && !r || !h && r)
+              continue;
+          }
+          const a = t.slice(e, e + this.blockSize), l = t.slice(e + 1, e + this.blockSize + 1);
+          yield { xs: a, ys: l };
+        }
     }).bind(this);
-    return d(k).batch(s).map((t) => {
-      const o = t;
-      return f(() => ({
-        xs: o.xs.cast("int32"),
-        ys: o.ys.cast("int32")
+    return S(n).batch(s).map((e) => {
+      const a = e;
+      return u(() => ({
+        xs: a.xs.cast("int32"),
+        ys: a.ys.cast("int32")
         // this.tf.oneHot(batchData.ys.cast('int32'), this.tokenizer.vocabSize),
       }));
     }).prefetch(2);
   }
 }
 export {
-  y as DatasetBuilder
+  w as DatasetBuilder,
+  p as PAGE_FACTOR,
+  y as flattenTokens
 };

package/dist/training/Trainer.js

@@ -1,13 +1,13 @@
-import { DatasetBuilder as d } from "./DatasetBuilder.js";
-import h from "./AdamExt.js";
-import { t as g, v as u, a as o } from "../index-CnHyhpKc.js";
-import { m as y, n as f } from "../norm-BpWsOapl.js";
-import { m as S, a as z } from "../moments-DLTE6-1p.js";
-import { m as b } from "../max-CcnEArWK.js";
-import { z as n } from "../zeros-CYMicyqz.js";
+import { DatasetBuilder as h, flattenTokens as d, PAGE_FACTOR as g } from "./DatasetBuilder.js";
+import u from "./AdamExt.js";
+import { t as f, v as y, a as m } from "../index-CnHyhpKc.js";
+import { m as S, n as z } from "../norm-BpWsOapl.js";
+import { m as w, a as T } from "../moments-DLTE6-1p.js";
+import { m as x } from "../max-CcnEArWK.js";
+import { z as p } from "../zeros-CYMicyqz.js";
 class G {
   constructor(t, s, e = 1e-3) {
-    this.tokenizer = s, this.model = t, this.learningRate = e, this.resetOptimizer(), this.datasetBuilder = new d(s, t.config.gpt.blockSize);
+    this.tokenizer = s, this.model = t, this.learningRate = e, this.resetOptimizer(), this.datasetBuilder = new h(s, t.config.gpt.blockSize);
   }
   model;
   optimizer;
@@ -29,7 +29,7 @@ class G {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const s = new h(
+    const s = new u(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -46,21 +46,21 @@ class G {
   printGradients(t) {
     Object.keys(t).forEach((s) => {
       const e = t[s];
-      console.log(`${s}:`), console.log(`  Shape: ${e.shape}`), console.log(`  Mean: ${S(e).dataSync()[0]}`), console.log(`  Std: ${z(e).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${y(e).dataSync()[0]}`), console.log(`  Max: ${b(e).dataSync()[0]}`), console.log(`  Norm: ${f(e).dataSync()[0]}`);
+      console.log(`${s}:`), console.log(`  Shape: ${e.shape}`), console.log(`  Mean: ${w(e).dataSync()[0]}`), console.log(`  Std: ${T(e).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${S(e).dataSync()[0]}`), console.log(`  Max: ${x(e).dataSync()[0]}`), console.log(`  Norm: ${z(e).dataSync()[0]}`);
     });
   }
   trainStep(t, s = !1, e = !1) {
-    return g(() => {
+    return f(() => {
       this.model.getProfiler()?.startMemory();
-      const { xs: a, ys: r } = t, l = () => {
-        const [m, p] = this.model.forward({ training: !0 }, a, r);
-        return m.dispose(), p;
-      }, { value: c, grads: i } = u(l);
-      return s ? this.model.getProfiler()?.endMemory("Training") : (e && (console.log("-------"), this.printGradients(i), console.log("-------")), this.optimizer.applyGradients(i), this.model.getProfiler()?.endMemory("Training"), o(i)), c;
+      const { xs: a, ys: i } = t, o = () => {
+        const [l, c] = this.model.forward({ training: !0 }, a, i);
+        return l.dispose(), c;
+      }, { value: n, grads: r } = y(o);
+      return s ? this.model.getProfiler()?.endMemory("Training") : (e && (console.log("-------"), this.printGradients(r), console.log("-------")), this.optimizer.applyGradients(r), this.model.getProfiler()?.endMemory("Training"), m(r)), n;
     });
   }
   dummyPass() {
-    const t = n([1, this.model.config.gpt.blockSize], "int32"), s = n([1, this.model.config.gpt.blockSize], "int32");
+    const t = p([1, this.model.config.gpt.blockSize], "int32"), s = p([1, this.model.config.gpt.blockSize], "int32");
     try {
       const e = this.trainStep({ xs: t, ys: s }, !0);
       e.dataSync(), e.dispose();
@@ -75,20 +75,29 @@ class G {
       const e = this.trainStep(s, !1, !1);
       return s.xs.dispose(), s.ys.dispose(), t.step++, t.totalSteps++, e.array().then((a) => (t.lastLoss = a, t.losses.push(t.lastLoss), e.dispose(), t.lastLoss));
     } catch (e) {
-      throw console.error(`Error processing batch at step ${t.step}:`, e), o(), e;
+      throw console.error(`Error processing batch at step ${t.step}:`, e), m(), e;
     }
   }
   async createTrainValidationSplit(t, s = 32, e = 0.1) {
-    const a = await this.datasetBuilder.createTextDataset(t, s, 0, 1 - e), r = await this.datasetBuilder.createTextDataset(
-      t,
+    const a = await d(t, this.tokenizer), i = /* @__PURE__ */ new Set();
+    if (e > 0) {
+      const r = Math.floor(a.length / (this.datasetBuilder.blockSize * g)), l = Math.max(1, Math.floor(r * e));
+      for (; i.size < l; ) {
+        const c = Math.floor(Math.random() * r);
+        i.add(c);
+      }
+    }
+    const o = await this.datasetBuilder.createTextDataset(a, s, i, !1), n = await this.datasetBuilder.createTextDataset(
+      a,
       s,
-      1 - e,
-      1
+      i,
+      !0
     );
-    return { trainDataset: a, validationDataset: r };
+    return { trainDataset: o, validationDataset: n };
   }
   async createDataset(t, s = 32) {
-    return await this.datasetBuilder.createTextDataset(t, s);
+    const e = await d(t, this.tokenizer);
+    return await this.datasetBuilder.createTextDataset(e, s);
   }
   dispose() {
     this.optimizer && this.optimizer.dispose();

package/dist/utilities/tokenParse.js

@@ -1,12 +1,21 @@
-function h(o) {
-  const r = Array.from(o), t = [], l = new RegExp("(\\p{P}|\\p{S}|\\s)", "gu");
-  let e = "";
-  for (let n = 0; n < r.length; n++) {
-    const s = r[n];
-    s === " " ? (t.push(e), e = s) : s.match(l) ? (t.push(e), t.push(s), e = "") : e += s;
+function c(l) {
+  const r = Array.from(l), s = [], o = new RegExp("(\\p{P}|\\p{S}|\\s)", "gu");
+  let t = "";
+  for (let e = 0; e < r.length; e++) {
+    const n = r[e];
+    if (n === " ")
+      (r[e + 1] ?? "") !== " " ? (s.push(t), t = n) : t += n;
+    else if (n.match(o)) {
+      s.push(t);
+      let h = n;
+      for (; e + 1 < r.length && r[e + 1] === n; )
+        h += r[e + 1], e++;
+      s.push(h), t = "";
+    } else
+      t += n;
   }
-  return e.length > 0 && t.push(e), t.filter((n) => n.length > 0);
+  return t.length > 0 && s.push(t), s.filter((e) => e.length > 0);
 }
 export {
-  h as default
+  c as default
 };

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.5.2",
+    "version": "0.5.3",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",