@epfml/discojs 2.1.2-p20240515133413.0 → 2.1.2-p20240531085945.0

package/dist/models/gpt/model.js

@@ -0,0 +1,191 @@
+import * as tf from '@tensorflow/tfjs';
+import { getModelSizes, DEFAULT_CONFIG } from './config.js';
+import { getCustomAdam, clipByGlobalNormObj } from './optimizers.js';
+import evaluate from './evaluate.js';
+import { GPTArchitecture } from './layers.js';
+/**
+ * GPTModel extends tf.LayersModel and overrides tfjs' default training loop
+ *
+ */
+class GPTModel extends tf.LayersModel {
+    config;
+    disposalRefs; // Array to store tensor to dispose manually
+    // Object to pass down to layers to store max memory allocated
+    // This is an object rather than a primitive to pass the reference
+    peakMemory;
+    constructor(partialConfig) {
+        // Fill missing config parameters with default values
+        let completeConfig = { ...DEFAULT_CONFIG, ...partialConfig };
+        // Add layer sizes depending on which model has been specified
+        completeConfig = { ...completeConfig, ...getModelSizes(completeConfig.modelType) };
+        // Init the tf.LayersModel and assign it to this
+        const disposalRefs = [];
+        const peakMemory = { value: 0 };
+        const gpt = GPTArchitecture(completeConfig, disposalRefs, peakMemory);
+        const { inputs, outputs, name } = gpt;
+        super({ inputs, outputs, name });
+        this.config = completeConfig;
+        this.disposalRefs = disposalRefs;
+        this.peakMemory = peakMemory;
+    }
+    // Some tensors are not cleaned up when model.dispose is called 
+    // So we dispose them manually
+    disposeRefs() {
+        for (const tensorContainer of this.disposalRefs) {
+            tf.dispose([tensorContainer]);
+        }
+    }
+    get getGPTConfig() {
+        return this.config;
+    }
+    compile() {
+        this.optimizer = this.config.weightDecay !== 0
+            ? getCustomAdam(this, this.config.lr, this.config.weightDecay)
+            : tf.train.adam(this.config.lr);
+        this.peakMemory.value = 0;
+    }
+    async fitDataset(dataset, trainingArgs) {
+        const callbacks = trainingArgs.callbacks;
+        const evalDataset = trainingArgs.validationData;
+        await callbacks.onTrainBegin?.();
+        for (let epoch = 1; epoch <= trainingArgs.epochs; epoch++) {
+            let averageLoss = 0;
+            let iteration = 1;
+            const iterator = await dataset.iterator();
+            let preprocessingTime = performance.now();
+            let next = await iterator.next();
+            preprocessingTime = performance.now() - preprocessingTime;
+            while (next.done !== true && iteration <= this.config.maxIter) {
+                let weightUpdateTime = performance.now();
+                await callbacks.onEpochBegin?.(epoch);
+                const { xs, ys } = next.value;
+                const lossFn = () => {
+                    const logits = this.apply(xs);
+                    if (Array.isArray(logits)) {
+                        throw new Error('model outputs too many tensor');
+                    }
+                    if (logits instanceof tf.SymbolicTensor) {
+                        throw new Error('model outputs symbolic tensor');
+                    }
+                    return tf.losses.softmaxCrossEntropy(ys, logits);
+                };
+                let backwardPassMemory = 0;
+                const lossTensor = tf.tidy(() => {
+                    const { grads, value: lossTensor } = this.optimizer.computeGradients(lossFn);
+                    const gradsClipped = clipByGlobalNormObj(grads, 1);
+                    this.optimizer.applyGradients(gradsClipped);
+                    backwardPassMemory = tf.memory().numBytes / 1024 / 1024 / 1024;
+                    return lossTensor;
+                });
+                const loss = await lossTensor.array();
+                averageLoss += loss;
+                weightUpdateTime = performance.now() - weightUpdateTime;
+                // Probably never the case. Empirically the attention mechanism always allocates
+                // more memory than the backward pass
+                if (backwardPassMemory > this.peakMemory.value) {
+                    this.peakMemory.value = backwardPassMemory;
+                }
+                tf.dispose([xs, ys, lossTensor]);
+                if (evalDataset !== undefined &&
+                    this.config.evaluateEvery !== undefined &&
+                    iteration % this.config.evaluateEvery == 0) {
+                    const iterationLogs = await evaluate(this, evalDataset, this.config.maxEvalBatches);
+                    console.log(iterationLogs);
+                }
+                console.log(`Epoch: ${epoch}`, `\tStep: ${iteration} / ${this.config.maxIter}`, `\tLoss: ${loss.toFixed(3)}`, `\tPeak memory: ${this.peakMemory.value.toFixed(2)} GB`, `\tNumber of tensors allocated: ${tf.memory().numTensors}`, `\tPreprocessing time: ${preprocessingTime.toFixed(0)} ms`, `\tWeight update time: ${weightUpdateTime.toFixed(0)} ms`);
+                iteration++;
+                next = await iterator.next();
+            }
+            // Memory leak: If we reached the last iteration rather than the end of the dataset, cleanup the tensors
+            if (next.done != true && iteration > this.config.maxIter) {
+                const { xs, ys } = next.value;
+                tf.dispose([xs, ys]);
+            }
+            let logs = {
+                'loss': averageLoss / iteration,
+                'peakMemory': this.peakMemory.value
+            };
+            if (evalDataset !== undefined) {
+                logs = { ...logs, ...await evaluate(this, evalDataset, this.config.maxEvalBatches) };
+                console.log(logs);
+            }
+            await callbacks.onEpochEnd?.(epoch, logs);
+        }
+        await callbacks.onTrainEnd?.();
+        return new tf.History();
+    }
+}
+const defaultGenerateConfig = {
+    maxNewTokens: 20,
+    temperature: 1.0,
+    doSample: false
+};
+function prepareIdx(idx) {
+    return tf.tidy(() => {
+        let ret;
+        if (idx instanceof tf.Tensor) {
+            ret = idx.clone();
+        }
+        else {
+            ret = tf.tensor(idx);
+        }
+        if (ret.dtype !== 'int32') {
+            ret = ret.toInt();
+        }
+        switch (ret.shape.length) {
+            case 1:
+                return ret.expandDims(0);
+            case 2:
+                return ret;
+            default:
+                throw new Error('unexpected shape');
+        }
+    });
+}
+/**
+ * GPTForCausalLM stands for GPT model for Causal Language Modeling. Causal because it only looks at past tokens and not future ones
+ * This class extends GPTModel and adds supports for text generation
+ *
+ */
+export class GPTForCausalLM extends GPTModel {
+    async generate(idxRaw, conf) {
+        const config = Object.assign({}, defaultGenerateConfig, conf);
+        let idx = prepareIdx(idxRaw);
+        for (let step = 0; step < config.maxNewTokens; step++) {
+            const idxNext = this.generateOnce(this, idx, config);
+            const idxNew = idx.concat(idxNext, 1);
+            tf.dispose(idx);
+            idx = idxNew;
+            tf.dispose(idxNext);
+        }
+        const idxArr = await idx.array();
+        tf.dispose(idx);
+        return idxArr;
+    }
+    generateOnce(model, idx, config) {
+        const idxNext = tf.tidy(() => {
+            // slice input tokens if longer than context length
+            const blockSize = this.config.blockSize;
+            idx = idx.shape[1] <= blockSize
+                ? idx : idx.slice([0, idx.shape[1] - blockSize]);
+            const output = model.predict(idx);
+            if (Array.isArray(output))
+                throw new Error('The model outputs too multiple values');
+            if (output.shape.length !== 3)
+                throw new Error('The model outputs wrong shape');
+            const logits = output;
+            const logitsScaled = logits
+                .slice([0, idx.shape[1] - 1, 0])
+                .reshape([logits.shape[0], logits.shape[2]])
+                .div(tf.scalar(config.temperature));
+            const probs = logitsScaled.softmax(-1);
+            if (config.doSample) {
+                return tf.multinomial(probs, 1);
+            }
+            else {
+                return probs.argMax(-1).expandDims(1);
+            }
+        });
+        return idxNext;
+    }
+}

package/dist/models/gpt/optimizers.d.ts

@@ -0,0 +1,4 @@
+import * as tf from '@tensorflow/tfjs';
+declare function clipByGlobalNormObj(tensorsObj: Record<string, tf.Tensor>, clipNorm: number, useNorm?: tf.Tensor): Record<string, tf.Tensor>;
+declare function getCustomAdam(model: tf.LayersModel, lr: number, weightDecay: number): tf.Optimizer;
+export { getCustomAdam, clipByGlobalNormObj };

package/dist/models/gpt/optimizers.js

@@ -0,0 +1,95 @@
+import * as tf from '@tensorflow/tfjs';
+function l2Loss(tensor) {
+    return tf.div(tf.sum(tf.square(tensor)), 2);
+}
+function globalNorm(tensors) {
+    const halfSquaredNorms = [];
+    tensors.forEach((tensor) => {
+        halfSquaredNorms.push(l2Loss(tensor));
+    });
+    const halfSquaredNorm = tf.sum(tf.stack(halfSquaredNorms));
+    const norm = tf.sqrt(tf.mul(halfSquaredNorm, tf.scalar(2.0, halfSquaredNorm.dtype)));
+    return norm;
+}
+function clipByGlobalNorm(tensors, clipNorm, useNorm) {
+    return tf.tidy(() => {
+        useNorm = useNorm ?? globalNorm(tensors);
+        const scale = tf.mul(clipNorm, tf.minimum(tf.div(tf.scalar(1.0), useNorm), tf.div(tf.scalar(1.0, useNorm.dtype), clipNorm)));
+        const tensorsClipped = [];
+        tensors.forEach((tensor) => {
+            tensorsClipped.push(tf.clone(tf.mul(tensor, scale)));
+        });
+        return tensorsClipped;
+    });
+}
+function clipByGlobalNormObj(tensorsObj, clipNorm, useNorm) {
+    const varNames = Object.keys(tensorsObj);
+    const tensorsArr = varNames.map((n) => tensorsObj[n]);
+    const tensorsArrClipped = clipByGlobalNorm(tensorsArr, clipNorm, useNorm);
+    const tensorsObjClipped = {};
+    tensorsArrClipped.forEach((t, ti) => {
+        tensorsObjClipped[varNames[ti]] = t;
+    });
+    return tensorsObjClipped;
+}
+class AdamW extends tf.AdamOptimizer {
+    weightDecayRate;
+    includeInWeightDecay;
+    excludeFromWeightDecay;
+    gradientClipNorm;
+    constructor(params) {
+        console.log('Using custom AdamW optimizer');
+        const defaultParams = {
+            learningRate: 0.1,
+            beta1: 0.9,
+            beta2: 0.999,
+            epsilon: 1e-7,
+            weightDecayRate: 0,
+            includeInWeightDecay: [],
+            excludeFromWeightDecay: [],
+            gradientClipNorm: 1.0
+        };
+        const p = Object.assign({}, defaultParams, params);
+        super(p.learningRate, p.beta1, p.beta2, p.epsilon);
+        this.weightDecayRate = p.weightDecayRate;
+        this.includeInWeightDecay = p.includeInWeightDecay;
+        this.excludeFromWeightDecay = p.excludeFromWeightDecay;
+        this.gradientClipNorm = p.gradientClipNorm;
+    }
+    applyGradients(variableGradients) {
+        const varNames = Array.isArray(variableGradients)
+            ? variableGradients.map((v) => v.name)
+            : Object.keys(variableGradients);
+        varNames.forEach((name) => {
+            if (this.includeInWeightDecay.includes(name)) {
+                const value = tf.engine().registeredVariables[name];
+                const newValue = tf.sub(value, tf.mul(this.learningRate, tf.mul(value, this.weightDecayRate)));
+                value.assign(newValue);
+            }
+        });
+        super.applyGradients(variableGradients);
+    }
+}
+function getCustomAdam(model, lr, weightDecay) {
+    const includeInWeightDecay = [];
+    const excludeFromWeightDecay = [];
+    // TODO unsafe cast
+    const namedWeights = model.getNamedWeights();
+    namedWeights.forEach((v) => {
+        if (v.name.includes('bias') ||
+            v.name.includes('normalization') ||
+            v.name.includes('emb')) {
+            excludeFromWeightDecay.push(v.name);
+        }
+        else {
+            includeInWeightDecay.push(v.name);
+        }
+    });
+    return new AdamW({
+        learningRate: lr,
+        weightDecayRate: weightDecay,
+        includeInWeightDecay,
+        excludeFromWeightDecay
+    });
+}
+export { getCustomAdam, clipByGlobalNormObj };

package/dist/models/index.d.ts

@@ -0,0 +1,5 @@
+export { EpochLogs, Model } from './model.js';
+export { GPT } from './gpt/index.js';
+export { GPTConfig } from './gpt/config.js';
+export { TFJS } from './tfjs.js';
+export { getTaskTokenizer } from './tokenizer.js';

package/dist/models/index.js

@@ -0,0 +1,4 @@
+export { Model } from './model.js';
+export { GPT } from './gpt/index.js';
+export { TFJS } from './tfjs.js';
+export { getTaskTokenizer } from './tokenizer.js';

package/dist/{default_tasks/simple_face/model.js → models/mobileNetV2_35_alpha_2_classes.js}

@@ -1,3 +1,5 @@
+// Source: https://storage.googleapis.com/deai-313515.appspot.com/models/mobileNetV2_35_alpha_2_classes/model.json
+// This model was converted using the tensorflow.js converter
 export default {
     format: "layers-model",
     generatedBy: "keras v2.6.0",

package/dist/{default_tasks/cifar10/model.js → models/mobileNet_v1_025_224.js}

@@ -1,3 +1,4 @@
+// Source: https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json
 export default {
     modelTopology: {
         keras_version: "2.1.4",

package/dist/models/model.d.ts

@@ -0,0 +1,51 @@
+/// <reference types="node" resolution-mode="require"/>
+import type tf from "@tensorflow/tfjs";
+import type { WeightsContainer } from "../index.js";
+import type { Dataset } from "../dataset/index.js";
+export interface EpochLogs {
+    epoch: number;
+    training: {
+        loss: number;
+        accuracy?: number;
+    };
+    validation?: {
+        loss: number;
+        accuracy: number;
+    };
+    peakMemory: number;
+}
+export type Prediction = tf.Tensor;
+export type Sample = tf.Tensor;
+/**
+ * Trainable predictor
+ *
+ * Allow for various implementation of models (various train function, tensor-library, ...)
+ **/
+export declare abstract class Model implements Disposable {
+    /** Return training state */
+    abstract get weights(): WeightsContainer;
+    /** Set training state */
+    abstract set weights(ws: WeightsContainer);
+    /**
+     * Improve predictor
+     *
+     * @param trainingData dataset to optimize for
+     * @param validationData dataset to measure how well it is training
+     * @param epochs number of pass over the training dataset
+     * @param tracker watch the various steps
+     * @yields on every epoch, training can be stop by `return`ing it
+     */
+    abstract train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs, void>;
+    /** Predict likely values */
+    abstract predict(input: Sample): Promise<Prediction>;
+    /**
+     * This method is automatically called to cleanup the memory occupied by the model
+     * when leaving the definition scope if the instance has been defined with the `using` keyword.
+     * For example:
+     * function f() {
+     *   using model = new Model();
+     * }
+     * Calling f() will call the model's dispose method when exiting the function.
+     */
+    abstract [Symbol.dispose](): void;
+}

package/dist/models/model.js

@@ -0,0 +1,8 @@
+/**
+ * Trainable predictor
+ *
+ * Allow for various implementation of models (various train function, tensor-library, ...)
+ **/
+// TODO make it typesafe: same shape of data/input/weights
+export class Model {
+}

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,22 @@
+import { AutoTokenizer } 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') {
+        tokenizer = await AutoTokenizer.from_pretrained(tokenizer);
+        task.trainingInformation.tokenizer = tokenizer;
+    }
+    return tokenizer;
+}

package/dist/validation/validator.js

@@ -21,14 +21,15 @@ export class Validator {
         }
     }
     async getLabel(ys) {
-        switch (ys.shape[1]) {
-            case 1:
-                return await ys.greaterEqual(tf.scalar(0.5)).data();
-            case 2:
-                return await ys.argMax(1).data();
-            default:
-                throw new Error(`unable to reduce tensor of shape: ${ys.shape.toString()}`);
+        // Binary classification
+        if (ys.shape[1] == 1) {
+            return await ys.greaterEqual(tf.scalar(0.5)).data();
+            // Multi-class classification
         }
+        else {
+            return await ys.argMax(-1).data();
+        }
+        // Multi-label classification is not supported
     }
     async assess(data, useConfusionMatrix = false) {
         const batchSize = this.task.trainingInformation?.batchSize;

package/package.json

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

package/dist/default_tasks/cifar10/index.d.ts

@@ -1,2 +0,0 @@
-import type { TaskProvider } from '../../index.js';
-export declare const cifar10: TaskProvider;

package/dist/default_tasks/simple_face/index.d.ts

@@ -1,2 +0,0 @@
-import type { TaskProvider } from '../../index.js';
-export declare const simpleFace: TaskProvider;