@epfml/discojs 2.1.2-p20240513140724.0 → 2.1.2-p20240515132210.0

package/dist/models/tfjs.d.ts

@@ -0,0 +1,24 @@
+import * as tf from '@tensorflow/tfjs';
+import { WeightsContainer } from '../index.js';
+import { Model } from './index.js';
+import type { EpochLogs, Prediction, Sample } from './model.js';
+import type { Dataset } from '../dataset/index.js';
+/** TensorFlow JavaScript model with standard training */
+export declare class TFJS extends Model {
+    private readonly model;
+    /** Wrap the given trainable model */
+    constructor(model: tf.LayersModel);
+    get weights(): WeightsContainer;
+    set weights(ws: WeightsContainer);
+    train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs>;
+    predict(input: Sample): Promise<Prediction>;
+    static deserialize(raw: tf.io.ModelArtifacts): Promise<Model>;
+    serialize(): Promise<tf.io.ModelArtifacts>;
+    [Symbol.dispose](): void;
+    /**
+     * extract wrapped model
+     *
+     * @deprecated use `Model` instead of relying on tf specifics
+     */
+    extract(): tf.LayersModel;
+}

package/dist/models/tfjs.js

@@ -0,0 +1,107 @@
+import * as tf from '@tensorflow/tfjs';
+import { WeightsContainer } from '../index.js';
+import { Model } from './index.js';
+/** TensorFlow JavaScript model with standard training */
+export class TFJS extends Model {
+    model;
+    /** Wrap the given trainable model */
+    constructor(model) {
+        super();
+        this.model = model;
+        if (model.loss === undefined) {
+            throw new Error('TFJS models need to be compiled to be used');
+        }
+    }
+    get weights() {
+        return new WeightsContainer(this.model.weights.map((w) => w.read()));
+    }
+    set weights(ws) {
+        this.model.setWeights(ws.weights);
+    }
+    async *train(trainingData, validationData, epochs = 1) {
+        for (let epoch = 0; epoch < epochs; epoch++) {
+            let logs;
+            let peakMemory = 0;
+            await this.model.fitDataset(trainingData, {
+                epochs: 1,
+                validationData,
+                callbacks: {
+                    onBatchEnd: (_) => {
+                        const currentMemory = tf.memory().numBytes / 1024 / 1024 / 1024; // GB
+                        if (currentMemory > peakMemory) {
+                            peakMemory = currentMemory;
+                        }
+                    },
+                    onEpochEnd: (_, cur) => { logs = cur; }
+                },
+            });
+            if (logs === undefined) {
+                throw new Error("Epoch didn't gave any logs");
+            }
+            const { loss, acc, val_acc, val_loss } = logs;
+            if (loss === undefined || isNaN(loss) || acc === undefined || isNaN(acc)) {
+                throw new Error("Training loss is undefined or nan");
+            }
+            const structuredLogs = {
+                epoch,
+                peakMemory,
+                training: {
+                    loss: logs.loss,
+                    accuracy: logs.acc,
+                }
+            };
+            if (validationData !== undefined) {
+                if (val_loss === undefined || isNaN(val_loss) ||
+                    val_acc === undefined || isNaN(val_acc)) {
+                    throw new Error("Invalid validation logs");
+                }
+                structuredLogs.validation = {
+                    accuracy: logs.val_acc,
+                    loss: logs.val_loss
+                };
+            }
+            yield structuredLogs;
+        }
+    }
+    predict(input) {
+        const ret = this.model.predict(input);
+        if (Array.isArray(ret)) {
+            throw new Error('prediction yield many Tensors but should have only returned one');
+        }
+        return Promise.resolve(ret);
+    }
+    static async deserialize(raw) {
+        return new this(await tf.loadLayersModel({
+            load: () => Promise.resolve(raw)
+        }));
+    }
+    async serialize() {
+        let resolveArtifacts;
+        const ret = new Promise((resolve) => { resolveArtifacts = resolve; });
+        await this.model.save({
+            save: (artifacts) => {
+                resolveArtifacts(artifacts);
+                return Promise.resolve({
+                    modelArtifactsInfo: {
+                        dateSaved: new Date(),
+                        modelTopologyType: 'JSON'
+                    }
+                });
+            }
+        }, {
+            includeOptimizer: true // keep model compiled
+        });
+        return await ret;
+    }
+    [Symbol.dispose]() {
+        this.model.dispose();
+    }
+    /**
+     * extract wrapped model
+     *
+     * @deprecated use `Model` instead of relying on tf specifics
+     */
+    extract() {
+        return this.model;
+    }
+}

package/dist/models/tokenizer.d.ts

@@ -0,0 +1,14 @@
+import type { Task } from '../index.js';
+import { PreTrainedTokenizer } from '@xenova/transformers';
+/**
+ * A task's tokenizer is initially specified as the tokenizer name, e.g., 'Xenova/gpt2'.
+ * The first time the tokenizer is needed, this function initializes the actual tokenizer object
+ * and saves it in the task' tokenizer field to be reused in subsequent calls.
+ *
+ * We are proceeding as such because the task object is sent from the server to the client. Rather than
+ * sending complex objects through the network, we simply send the tokenizer name, which is then initialized client-side the
+ * first time it is called.
+ * @param task the task object specifying which tokenizer to use
+ * @returns an initialized tokenizer object
+ */
+export declare function getTaskTokenizer(task: Task): Promise<PreTrainedTokenizer>;

package/dist/models/tokenizer.js

@@ -0,0 +1,23 @@
+import { AutoTokenizer, env } from '@xenova/transformers';
+/**
+ * A task's tokenizer is initially specified as the tokenizer name, e.g., 'Xenova/gpt2'.
+ * The first time the tokenizer is needed, this function initializes the actual tokenizer object
+ * and saves it in the task' tokenizer field to be reused in subsequent calls.
+ *
+ * We are proceeding as such because the task object is sent from the server to the client. Rather than
+ * sending complex objects through the network, we simply send the tokenizer name, which is then initialized client-side the
+ * first time it is called.
+ * @param task the task object specifying which tokenizer to use
+ * @returns an initialized tokenizer object
+ */
+export async function getTaskTokenizer(task) {
+    let tokenizer = task.trainingInformation.tokenizer;
+    if (tokenizer === undefined)
+        throw Error('No tokenizer specified in the task training information');
+    if (typeof tokenizer == 'string') {
+        env.allowLocalModels = false;
+        tokenizer = await AutoTokenizer.from_pretrained(tokenizer);
+        task.trainingInformation.tokenizer = tokenizer;
+    }
+    return tokenizer;
+}

package/dist/training/trainer/trainer_builder.js

@@ -1,4 +1,4 @@
-import { ModelType } from '../../index.js';
+import { StoredModelType } from '../../index.js';
 import { DistributedTrainer } from './distributed_trainer.js';
 import { LocalTrainer } from './local_trainer.js';
 /**
@@ -36,7 +36,7 @@ export class TrainerBuilder {
         if (modelID === undefined) {
             throw new TypeError('model ID is undefined');
         }
-        const info = { type: ModelType.WORKING, taskID: this.task.id, name: modelID };
+        const info = { type: StoredModelType.WORKING, taskID: this.task.id, name: modelID };
         const model = await (await this.memory.contains(info) ? this.memory.getModel(info) : client.getLatestModel());
         return model;
     }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epfml/discojs",
-  "version": "2.1.2-p20240513140724.0",
+  "version": "2.1.2-p20240515132210.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",