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

package/dist/models/gpt/layers.js

@@ -0,0 +1,272 @@
+import * as tf from '@tensorflow/tfjs';
+/**
+ * Defines a range, from 0 to T, that is used to create positional embeddings
+ */
+class Range extends tf.layers.Layer {
+    static className = 'Range';
+    computeOutputShape(inputShape) {
+        return inputShape;
+    }
+    call(input, kwargs) {
+        return tf.tidy(() => {
+            if (Array.isArray(input)) {
+                // TODO support multitensor
+                input = input[0];
+            }
+            this.invokeCallHook(input, kwargs);
+            const T = input.shape[1];
+            if (T === undefined)
+                throw new Error('unexpected shape');
+            return tf.reshape(tf.range(0, T, 1, 'int32'), [1, T]);
+        });
+    }
+}
+tf.serialization.registerClass(Range);
+class LogLayer extends tf.layers.Layer {
+    static className = 'LogLayer';
+    computeOutputShape(inputShape) {
+        return inputShape;
+    }
+    call(input, kwargs) {
+        return tf.tidy(() => {
+            if (Array.isArray(input)) {
+                input = input[0];
+            }
+            this.invokeCallHook(input, kwargs);
+            return input;
+        });
+    }
+}
+tf.serialization.registerClass(LogLayer);
+class CausalSelfAttention extends tf.layers.Layer {
+    config;
+    peakMemory;
+    static className = 'CausalSelfAttention';
+    nHead;
+    nEmbd;
+    dropout;
+    bias;
+    mask;
+    cAttnKernel;
+    cAttnBias;
+    cProjKernel;
+    cProjBias;
+    constructor(config, disposalRefs, peakMemory) {
+        super(config);
+        this.config = config;
+        this.peakMemory = peakMemory;
+        this.nEmbd = config.nEmbd;
+        this.nHead = config.nHead;
+        this.dropout = config.dropout;
+        this.bias = config.bias;
+        // mask is a lower triangular matrix filled with 1
+        // calling bandPart zero out the upper triangular part of the all-ones matrix
+        // from the doc: tf.linalg.band_part(input, -1, 0) ==> Lower triangular part
+        this.mask = tf.linalg.bandPart(tf.ones([config.blockSize, config.blockSize]), -1, 0);
+        disposalRefs.push(this.mask); // Push a reference to dispose this matrix later
+    }
+    build() {
+        this.cAttnKernel = this.addWeight('c_attn/kernel', [this.nEmbd, 3 * this.nEmbd], 'float32', tf.initializers.glorotNormal({}));
+        this.cAttnBias = this.addWeight('c_attn/bias', [3 * this.nEmbd], 'float32', tf.initializers.zeros());
+        this.cProjKernel = this.addWeight('c_proj/kernel', [this.nEmbd, this.nEmbd], 'float32', tf.initializers.glorotNormal({}));
+        this.cProjBias = this.addWeight('c_proj/bias', [this.nEmbd], 'float32', tf.initializers.zeros());
+    }
+    computeOutputShape(inputShape) {
+        return inputShape;
+    }
+    getConfig() {
+        const config = super.getConfig();
+        return Object.assign({}, config, this.config);
+    }
+    call(input, kwargs) {
+        return tf.tidy(() => {
+            if (this.cAttnKernel === undefined ||
+                this.cAttnBias === undefined ||
+                this.cProjKernel === undefined ||
+                this.cProjBias === undefined) {
+                throw new Error('not built');
+            }
+            if (Array.isArray(input)) {
+                input = input[0];
+            }
+            this.invokeCallHook(input, kwargs);
+            const dense = (x, kernel, bias) => {
+                const k = kernel.read().expandDims(0).tile([x.shape[0], 1, 1]);
+                const m = x.matMul(k);
+                if (this.bias) {
+                    return tf.add(m, bias.read());
+                }
+                else {
+                    return m;
+                }
+            };
+            // Apply attention weights to inputs as one big matrix which is then split into the
+            // query, key and value submatrices
+            const cAttn = dense(input, this.cAttnKernel, this.cAttnBias);
+            let [q, k, v] = tf.split(cAttn, 3, -1);
+            const [B, T, C] = k.shape;
+            const splitHeads = (x) => tf.transpose(tf.reshape(x, [B, T, this.nHead, C / this.nHead]), [0, 2, 1, 3]);
+            q = splitHeads(q);
+            k = splitHeads(k);
+            v = splitHeads(v);
+            // Scaled self attention: query @ key / sqrt(n_heads)
+            let att = tf.mul(tf.matMul(q, k, false, true), tf.div(1, tf.sqrt(tf.cast(k.shape[k.shape.length - 1], 'float32'))));
+            // The next operations apply attention to the past tokens, which is
+            // essentially a weighted average of the past tokens with complicated weights,
+            // and makes sure to not pay any attention to future tokens
+            // mask is lower triangular matrix filled with 1
+            const mask = this.mask.slice([0, 0], [T, T]);
+            // 1 - mask                   => upper triangular matrix filled with 1
+            // (1 - mask) * -10^9         => upper triangular matrix filled with -inf
+            // att + ((1 - mask) * -10^9) => lower triangular part is the same as the `att` matrix
+            //                               upper triangular part is -inf
+            att = tf.add(att, tf.mul(tf.sub(1, mask), -1e9));
+            // applying softmax zeros out the upper triangular part 
+            //(which are the attention weights of future tokens)
+            // and creates a probability distribution for the lower triangular
+            // (attention weights of past tokens). The probability distribution ensures
+            // that the attention weights of past tokens for a particular token sum to one
+            att = tf.softmax(att, -1);
+            att = kwargs.training === true ? tf.dropout(att, this.dropout) : att;
+            // This is where the (attention-)weighted sum of past values is performed
+            let y = tf.matMul(att, v);
+            y = tf.transpose(y, [0, 2, 1, 3]);
+            y = tf.reshape(y, [B, T, C]);
+            y = dense(y, this.cProjKernel, this.cProjBias);
+            y = kwargs.training === true ? tf.dropout(y, this.dropout) : y;
+            const memoryAllocated = tf.memory().numBytes / 1024 / 1024 / 1024; // GB
+            if (memoryAllocated > this.peakMemory.value) {
+                this.peakMemory.value = memoryAllocated;
+            }
+            return y;
+        });
+    }
+}
+tf.serialization.registerClass(CausalSelfAttention);
+class GELU extends tf.layers.Layer {
+    static className = 'GELU';
+    constructor() {
+        super({});
+    }
+    computeOutputShape(inputShape) {
+        return inputShape;
+    }
+    call(input, kwargs) {
+        return tf.tidy(() => {
+            if (Array.isArray(input)) {
+                // TODO support multitensor
+                input = input[0];
+            }
+            this.invokeCallHook(input, kwargs);
+            const cdf = tf.mul(0.5, tf.add(1, tf.tanh(tf.mul(tf.sqrt(tf.div(2, Math.PI)), tf.add(input, tf.mul(0.044715, tf.pow(input, 3)))))));
+            return tf.mul(input, cdf);
+        });
+    }
+}
+tf.serialization.registerClass(GELU);
+function MLP(config) {
+    return tf.sequential({ layers: [
+            tf.layers.dense({
+                name: 'mlp/c_fc',
+                units: 4 * config.nEmbd,
+                inputDim: config.nEmbd,
+                inputShape: [config.blockSize, config.nEmbd]
+            }),
+            new GELU(),
+            tf.layers.dense({
+                name: 'mlp/c_proj',
+                units: config.nEmbd,
+                inputDim: 4 * config.nEmbd,
+                inputShape: [config.blockSize, 4 * config.nEmbd]
+            }),
+            tf.layers.dropout({
+                name: 'mlp/drop',
+                rate: config.residDrop
+            }),
+        ] });
+}
+function TransformerBlock(conf, disposalRefs, peakMemory) {
+    const config = Object.assign({ name: 'h' }, conf);
+    const inputs = tf.input({ shape: [config.blockSize, config.nEmbd] });
+    let x1, x2;
+    // input normalization
+    x1 = tf.layers.layerNormalization({ name: config.name + '/ln_1', epsilon: 1e-5 })
+        .apply(inputs);
+    if (config.debug) {
+        x1 = new LogLayer({ name: config.name + '/ln_1_log' }).apply(x1);
+    }
+    // self attention layer
+    x1 = new CausalSelfAttention(Object.assign({}, config, { name: config.name + '/attn' }), disposalRefs, peakMemory).apply(x1);
+    // Residual connection
+    x1 = tf.layers.add().apply([inputs, x1]);
+    // normalization 
+    x2 = tf.layers
+        .layerNormalization({ name: config.name + '/ln_2', epsilon: 1e-5 })
+        .apply(x1);
+    // MLP 
+    x2 = MLP(Object.assign({}, config, { name: config.name + '/mlp' })).apply(x2);
+    // add attention output to mlp output
+    x2 = tf.layers.add().apply([x1, x2]);
+    return tf.model({ name: config.name, inputs, outputs: x2 });
+}
+/**
+ * The GPTArchitecture specifically defines a GPT forward pass, i.e.,
+ * what are the inputs, the successive transformer blocks and the outputs. It is then
+ * used to create a GPTModel
+ *
+ * @param conf GPTConfig
+ * @returns model, tf.LayersModel, which supports model(inputs), model.predict and model.apply
+ */
+export function GPTArchitecture(config, disposalRefs, peakMemory) {
+    const inputs = tf.input({ shape: [null] });
+    //Token embedding
+    const tokEmb = config.tokEmb
+        ? tf.layers.embedding({
+            name: config.name + '/wte',
+            inputDim: config.vocabSize,
+            outputDim: config.nEmbd,
+            embeddingsInitializer: 'zeros',
+            embeddingsRegularizer: undefined,
+            activityRegularizer: undefined
+        }).apply(inputs)
+        : inputs;
+    // Positional embedding
+    const range = new Range({}).apply(inputs);
+    let posEmb = tf.layers.embedding({
+        name: config.name + '/wpe',
+        inputDim: config.blockSize,
+        outputDim: config.nEmbd,
+        embeddingsInitializer: 'zeros'
+    }).apply(range);
+    if (config.debug) {
+        posEmb = new LogLayer({ name: 'posEmb' }).apply(posEmb);
+    }
+    // token and positional embeddings are added together
+    let x = tf.layers.add().apply([tokEmb, posEmb]);
+    // dropout
+    x = tf.layers.dropout({ name: 'drop', rate: config.embdDrop }).apply(x);
+    if (config.debug) {
+        x = new LogLayer({ name: 'dropadd' }).apply(x);
+    }
+    //Apply successively transformer blocks, attention and dense layers
+    for (let i = 0; i < config.nLayer; i++) {
+        x = TransformerBlock(Object.assign({}, config, { name: config.name + '/h/' + i }), disposalRefs, peakMemory).apply(x);
+    }
+    // Normalization
+    x = tf.layers.layerNormalization({ name: config.name + '/ln_f', epsilon: 1e-5 })
+        .apply(x);
+    if (config.debug) {
+        x = new LogLayer({ name: 'fin/ln' }).apply(x);
+    }
+    // Append a language modeling head if specified
+    if (config.lmHead) {
+        x = tf.layers.dense({
+            name: 'lm_head',
+            units: config.vocabSize,
+            inputDim: config.nEmbd,
+            inputShape: [config.blockSize, config.nEmbd],
+            useBias: false
+        }).apply(x);
+    }
+    return tf.model({ inputs, outputs: x });
+}

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

@@ -0,0 +1,43 @@
+import * as tf from '@tensorflow/tfjs';
+import type { GPTConfig } from './config.js';
+/**
+ * tfjs does not export LazyIterator and Dataset...
+ */
+declare abstract class LazyIterator<T> {
+    abstract next(): Promise<IteratorResult<T>>;
+}
+export declare abstract class Dataset<T> {
+    abstract iterator(): Promise<LazyIterator<T>>;
+    size: number;
+}
+/**
+ * GPTModel extends tf.LayersModel and overrides tfjs' default training loop
+ *
+ */
+declare class GPTModel extends tf.LayersModel {
+    protected readonly config: Required<GPTConfig>;
+    private readonly disposalRefs;
+    protected peakMemory: {
+        value: number;
+    };
+    constructor(partialConfig?: GPTConfig);
+    disposeRefs(): void;
+    get getGPTConfig(): Required<GPTConfig>;
+    compile(): void;
+    fitDataset<T>(dataset: Dataset<T>, trainingArgs: tf.ModelFitDatasetArgs<T>): Promise<tf.History>;
+}
+interface GenerateConfig {
+    maxNewTokens: number;
+    temperature: number;
+    doSample: boolean;
+}
+/**
+ * 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 declare class GPTForCausalLM extends GPTModel {
+    generate(idxRaw: tf.TensorLike, conf: GenerateConfig): Promise<number[][]>;
+    private generateOnce;
+}
+export {};

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;
+}