@genai-fi/nanogpt 1.2.4 → 1.3.0

package/dist/api/training.d.ts

@@ -64,7 +64,7 @@ export default class Training {
     getJob(id: string): ITrainingJob | null;
     /** Resume a paused job and optionally change some options. */
     resume(id: string): Promise<void>;
-    cancel(id: string): void;
+    cancel(id?: string): void;
     addBreak(jobid: string): number | undefined;
     deleteBreak(jobid: string, breakId: number): void;
     getPretrainingJob(): ITrainingJob | null;

package/dist/api/training.js

@@ -108,7 +108,7 @@ var o = class {
 		return s.trainer.log = t, s.trainer.resumeFromLog(t[t.length - 1]), this._jobs.set(s.id, s), s;
 	}
 	launchJob(e) {
-		e.trainDataset && (this.setState(e, "running"), e.trainer.trainOnDataset(e.trainDataset, e.options, e.validationDataset, (t) => {
+		e.trainDataset && (this.setState(e, "running"), e.trainer.prepareToTrain(e.options), e.trainer.trainOnDataset(e.trainDataset, e.options, e.validationDataset, (t) => {
 			e.history = e.trainer.log, e.progress = t.totalTokens / e.totalTokens, e.remaining = Math.max(0, (e.totalTokens - t.totalTokens) / t.totalTokens * t.duration), e.breakOnLog.size > 0 && e.state === "running" && (this.setState(e, "pausing"), e.trainer.stop()), this.ee.emit("progress", e);
 		}).then(() => {
 			e.state === "running" ? this.setState(e, "completed") : e.state === "pausing" ? this.setState(e, "paused") : e.state === "cancelling" && this.setState(e, "cancelled");
@@ -180,12 +180,15 @@ var o = class {
 		t && t.state === "paused" && (t.resumeCounter++, !(t.resumeCounter < t.breakOnLog.size) && (t.resumeCounter = 0, this.setState(t, "pending"), this.launchJob(t)));
 	}
 	cancel(e) {
-		this.assertValidJobId(e);
-		let t = this.getJob(e);
-		if (t) if (t.state === "paused") {
-			this.setState(t, "cancelled");
-			return;
-		} else (t.state === "running" || t.state === "pausing") && (this.setState(t, "cancelling"), t.trainer.stop());
+		if (e) {
+			this.assertValidJobId(e);
+			let t = this.getJob(e);
+			if (!t) return;
+			if (t.state === "paused") {
+				this.setState(t, "cancelled");
+				return;
+			} else (t.state === "running" || t.state === "pausing") && (this.setState(t, "cancelling"), t.trainer.stop());
+		} else for (let e of this._jobs.values()) e.state === "paused" ? this.setState(e, "cancelled") : (e.state === "running" || e.state === "pausing") && (this.setState(e, "cancelling"), e.trainer.stop());
 	}
 	addBreak(e) {
 		let t = this.getJob(e);

package/dist/training/BasicTrainer.d.ts

@@ -45,6 +45,13 @@ export default class BasicTrainer {
     dispose(): void;
     private createEmptyState;
     private performLogging;
+    /** A sync function to ensure the trainer is prepared and the state is initialized before training.
+     *  It should be called before invoking trainOnDataset.
+     */
+    prepareToTrain(options: Partial<TrainingOptions>): {
+        losses: number[];
+        validationLosses: number[];
+    } | undefined;
     trainOnDataset(dataset: Dataset<{
         xs: Tensor;
         ys: Tensor;

package/dist/training/BasicTrainer.js

@@ -211,22 +211,31 @@ var p = {
 		}
 		this.log.push(p), r && r(p), a.logStartTime = Date.now();
 	}
-	async trainOnDataset(e, t, n, i) {
-		let { logInterval: a = 40, maxEpochs: o = Infinity } = {
+	prepareToTrain(e) {
+		let { maxEpochs: t = Infinity } = {
 			...p,
-			...t
+			...e
 		};
 		this.log.length > 0 && this.resumeFromLog(this.log[this.log.length - 1]);
-		let c = o * (t?.epochSteps || 1e3);
-		t.metrics && this.setMetrics(t.metrics);
-		let l = this.createEmptyState();
-		if (this.lastState = l, l.step >= c) return {
-			losses: l.losses,
-			validationLosses: l.validationLosses
+		let n = t * (e?.epochSteps || 1e3);
+		e.metrics && this.setMetrics(e.metrics);
+		let r = this.createEmptyState();
+		if (this.lastState = r, r.step >= n) return {
+			losses: r.losses,
+			validationLosses: r.validationLosses
 		};
+		this.running = !0;
+	}
+	async trainOnDataset(e, t, n, i) {
+		let { logInterval: a = 40, maxEpochs: o = Infinity } = {
+			...p,
+			...t
+		}, c = o * (t?.epochSteps || 1e3);
 		await this.dummyPass(), t?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new s()));
+		let l = this.lastState;
+		if (!l) throw Error("Trainer state has not been initialized. Call prepareToTrain() first.");
 		let d = Date.now();
-		this.running = !0, l.logStartTime = d;
+		l.logStartTime = d;
 		let f = d, m = n ? new u(this.model, n, this.maskedLoss) : void 0, h = await e.iterator(), g = h.next();
 		try {
 			for (; this.running;) {

package/package.json

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