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

package/dist/default_tasks/skin_condition.js

@@ -0,0 +1,79 @@
+import * as tf from '@tensorflow/tfjs';
+import { data, models } from '../index.js';
+const IMAGE_SIZE = 128;
+const LABELS = ['Eczema', 'Allergic Contact Dermatitis', 'Urticaria'];
+export const skinCondition = {
+    getTask() {
+        return {
+            id: 'skin_condition',
+            displayInformation: {
+                taskTitle: 'Skin Condition Classification',
+                summary: {
+                    preview: "Identify common skin conditions from volunteer image contributions. You can find a sample dataset of 400 images <a class='underline text-primary-dark dark:text-primary-light' href='https://storage.googleapis.com/deai-313515.appspot.com/scin_sample.zip'>here</a> or see the full <a class='underline text-primary-dark dark:text-primary-light' href='https://github.com/google-research-datasets/scin/tree/main'>SCIN dataset</a>. You can find how to download and preprocess the dataset <a class='underline text-primary-dark dark:text-primary-light' href='https://github.com/epfml/disco/blob/develop/docs/examples/scin_dataset.ipynb'>in this notebook</a>.",
+                    overview: "The <a class='underline text-primary-dark dark:text-primary-light' href='https://github.com/google-research-datasets/scin/tree/main'>SCIN (Skin Condition Image Network) open access dataset</a> aims to supplement publicly available dermatology datasets from health system sources with representative images from internet users. To this end, the SCIN dataset was collected from Google Search users in the United States through a voluntary, consented image donation application. The SCIN dataset is intended for health education and research, and to increase the diversity of dermatology images available for public use. The SCIN dataset contains 5,000+ volunteer contributions (10,000+ images) of common dermatology conditions. Contributions include Images, self-reported demographic, history, and symptom information, and self-reported Fitzpatrick skin type (sFST). In addition, dermatologist labels of the skin condition are provided for each contribution. You can find more information on the dataset and classification task <a class='underline text-primary-dark dark:text-primary-light' href='https://arxiv.org/abs/2402.18545'>here</a>."
+                },
+                dataFormatInformation: "There are hundreds of skin condition labels in the SCIN dataset. For the sake of simplicity, we only include the 3 most common conditions in the sample dataset: 'Eczema', 'Allergic Contact Dermatitis' and 'Urticaria'. Therefore, each image is expected to be labeled with one of these three categories.",
+                sampleDatasetLink: 'https://storage.googleapis.com/deai-313515.appspot.com/scin_sample.zip'
+            },
+            trainingInformation: {
+                modelID: 'skin-condition-model',
+                epochs: 10,
+                roundDuration: 2,
+                validationSplit: 0.3,
+                batchSize: 8,
+                preprocessingFunctions: [data.ImagePreprocessing.Resize, data.ImagePreprocessing.Normalize],
+                dataType: 'image',
+                IMAGE_H: IMAGE_SIZE,
+                IMAGE_W: IMAGE_SIZE,
+                LABEL_LIST: LABELS,
+                scheme: 'federated',
+                noiseScale: undefined,
+                clippingRadius: undefined
+            }
+        };
+    },
+    async getModel() {
+        const imageChannels = 3;
+        const numOutputClasses = LABELS.length;
+        const model = tf.sequential();
+        model.add(tf.layers.conv2d({
+            inputShape: [IMAGE_SIZE, IMAGE_SIZE, imageChannels],
+            filters: 8,
+            kernelSize: 3,
+            strides: 1,
+            kernelInitializer: 'varianceScaling',
+            activation: 'relu'
+        }));
+        model.add(tf.layers.maxPooling2d({ poolSize: [2, 2] }));
+        model.add(tf.layers.dropout({ rate: 0.2 }));
+        const convFilters = [16, 32, 64, 128];
+        for (const filters of convFilters) {
+            model.add(tf.layers.conv2d({
+                filters: filters,
+                kernelSize: 3,
+                strides: 1,
+                kernelInitializer: 'varianceScaling',
+                activation: 'relu'
+            }));
+            model.add(tf.layers.maxPooling2d({ poolSize: [2, 2] }));
+            model.add(tf.layers.dropout({ rate: 0.2 }));
+        }
+        model.add(tf.layers.flatten());
+        model.add(tf.layers.dense({
+            units: 64,
+            kernelInitializer: 'varianceScaling',
+            activation: 'relu',
+        }));
+        model.add(tf.layers.dense({
+            units: numOutputClasses,
+            kernelInitializer: 'varianceScaling',
+            activation: 'softmax'
+        }));
+        model.compile({
+            optimizer: tf.train.adam(),
+            loss: 'categoricalCrossentropy',
+            metrics: ['accuracy']
+        });
+        return Promise.resolve(new models.TFJS(model));
+    }
+};

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

@@ -0,0 +1,32 @@
+type ModelType = 'gpt2' | 'gpt2-medium' | 'gpt2-large' | 'gpt2-xl' | 'gpt-mini' | 'gpt-micro' | 'gpt-nano';
+export interface GPTConfig {
+    lr: number;
+    blockSize: number;
+    vocabSize: number;
+    modelType: ModelType;
+    name?: string;
+    evaluate?: boolean;
+    maxEvalBatches?: number;
+    evaluateEvery?: number;
+    maxIter?: number;
+    weightDecay?: number;
+    verbose?: 0 | 1;
+    bias?: boolean;
+    debug?: boolean;
+    dropout?: number;
+    residDrop?: number;
+    embdDrop?: number;
+    tokEmb?: boolean;
+    lmHead?: boolean;
+    nLayer?: number;
+    nHead?: number;
+    nEmbd?: number;
+}
+export declare const DEFAULT_CONFIG: Required<GPTConfig>;
+export type ModelSize = {
+    nLayer: number;
+    nHead: number;
+    nEmbd: number;
+};
+export declare function getModelSizes(modelType: ModelType): Required<ModelSize>;
+export {};

package/dist/models/gpt/config.js

@@ -0,0 +1,42 @@
+// for a benchmark of performance, see https://github.com/epfml/disco/pull/659
+export const DEFAULT_CONFIG = {
+    name: 'transformer',
+    lr: 0.001,
+    weightDecay: 0,
+    maxIter: 5,
+    verbose: 0,
+    modelType: 'gpt-nano',
+    evaluate: true,
+    maxEvalBatches: 12,
+    evaluateEvery: 100,
+    blockSize: 128,
+    vocabSize: 50258,
+    bias: true,
+    debug: false,
+    dropout: 0.2,
+    residDrop: 0.2,
+    embdDrop: 0.2,
+    tokEmb: true,
+    lmHead: true,
+    nLayer: 3,
+    nHead: 3,
+    nEmbd: 48,
+};
+export function getModelSizes(modelType) {
+    switch (modelType) {
+        case 'gpt2':
+            return { nLayer: 12, nHead: 12, nEmbd: 768 };
+        case 'gpt2-medium':
+            return { nLayer: 24, nHead: 16, nEmbd: 1024 };
+        case 'gpt2-large':
+            return { nLayer: 36, nHead: 20, nEmbd: 1280 };
+        case 'gpt2-xl':
+            return { nLayer: 48, nHead: 25, nEmbd: 1600 };
+        case 'gpt-mini':
+            return { nLayer: 6, nHead: 6, nEmbd: 192 };
+        case 'gpt-micro':
+            return { nLayer: 4, nHead: 4, nEmbd: 128 };
+        case 'gpt-nano':
+            return { nLayer: 3, nHead: 3, nEmbd: 48 };
+    }
+}

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

@@ -0,0 +1,7 @@
+import * as tf from '@tensorflow/tfjs';
+interface DataPoint extends tf.TensorContainerObject {
+    xs: tf.Tensor2D;
+    ys: tf.Tensor3D;
+}
+export default function evaluate(model: tf.LayersModel, dataset: tf.data.Dataset<DataPoint>, maxEvalBatches: number): Promise<Record<'acc' | 'val_acc' | 'val_loss' | 'val_perplexity', number>>;
+export {};

package/dist/models/gpt/evaluate.js

@@ -0,0 +1,44 @@
+import * as tf from '@tensorflow/tfjs';
+export default async function evaluate(model, dataset, maxEvalBatches) {
+    let datasetSize = 0;
+    let totalLoss = 0;
+    const acc = [0, 0];
+    await dataset.take(maxEvalBatches).map(({ xs, ys }) => {
+        const logits = model.apply(xs);
+        if (Array.isArray(logits)) {
+            throw new Error('model output too many tensor');
+        }
+        if (logits instanceof tf.SymbolicTensor) {
+            throw new Error('model output symbolic tensor');
+        }
+        xs.dispose();
+        return { logits, ys };
+    }).mapAsync(async ({ logits, ys }) => {
+        const lossTensor = tf.losses.softmaxCrossEntropy(ys, logits);
+        const loss = await lossTensor.array();
+        if (typeof loss !== 'number') {
+            throw new Error('got multiple loss');
+        }
+        const accTensor = tf.metrics.categoricalAccuracy(ys, logits);
+        const accSize = accTensor.shape.reduce((l, r) => l * r, 1);
+        const accSum = accTensor.sum();
+        const accSummed = await accSum.array();
+        if (typeof accSummed !== 'number') {
+            throw new Error('got multiple accuracy sum');
+        }
+        tf.dispose([ys, logits, accTensor, accSum, lossTensor]);
+        return { loss, accSummed, accSize };
+    }).forEachAsync(({ loss, accSummed, accSize }) => {
+        datasetSize += 1;
+        totalLoss += loss;
+        acc[0] += accSummed;
+        acc[1] += accSize;
+    });
+    const loss = totalLoss / datasetSize;
+    return {
+        val_loss: loss,
+        val_perplexity: Math.exp(loss),
+        acc: acc[0] / acc[1],
+        val_acc: acc[0] / acc[1]
+    };
+}

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

@@ -0,0 +1,35 @@
+/**
+ * this code is taken from gpt-tfjs with modifications from @peacefulotter and @lukemovement
+ **/
+import { PreTrainedTokenizer } from '@xenova/transformers';
+import { WeightsContainer } from '../../index.js';
+import type { Dataset } from '../../dataset/index.js';
+import { Model } from '../model.js';
+import type { EpochLogs, Prediction, Sample } from '../model.js';
+import type { GPTConfig } from './config.js';
+export declare class GPT extends Model {
+    private readonly model;
+    constructor(partialConfig?: GPTConfig);
+    /**
+     * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs)
+     * This allows for getting logs and stopping training without callbacks.
+     *
+     * @param trainingData training dataset
+     * @param validationData validation dataset
+     * @param epochs the number of passes of the training dataset
+     * @param tracker
+     */
+    train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs, void>;
+    predict(input: Sample): Promise<Prediction>;
+    generate(input: string, tokenizer: PreTrainedTokenizer, newTokens?: number): Promise<string>;
+    get config(): Required<GPTConfig>;
+    get weights(): WeightsContainer;
+    set weights(ws: WeightsContainer);
+    static deserialize(data: GPTSerialization): Model;
+    serialize(): GPTSerialization;
+    [Symbol.dispose](): void;
+}
+export type GPTSerialization = {
+    weights: WeightsContainer;
+    config?: GPTConfig;
+};

package/dist/models/gpt/index.js

@@ -0,0 +1,104 @@
+/**
+ * this code is taken from gpt-tfjs with modifications from @peacefulotter and @lukemovement
+ **/
+import { WeightsContainer } from '../../index.js';
+import { Model } from '../model.js';
+import { GPTForCausalLM } from './model.js';
+export class GPT extends Model {
+    model;
+    constructor(partialConfig) {
+        super();
+        this.model = new GPTForCausalLM(partialConfig);
+    }
+    /**
+     * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs)
+     * This allows for getting logs and stopping training without callbacks.
+     *
+     * @param trainingData training dataset
+     * @param validationData validation dataset
+     * @param epochs the number of passes of the training dataset
+     * @param tracker
+     */
+    async *train(trainingData, validationData, epochs = 1) {
+        this.model.compile();
+        let logs;
+        const trainingArgs = {
+            epochs: 1, // force fitDataset to do only one epoch because it is wrapped in a for loop
+            validationData,
+            callbacks: { onEpochEnd: (_, cur) => { logs = cur; } },
+        };
+        for (let epoch = 0; epoch < epochs; epoch++) {
+            await this.model.fitDataset(trainingData, trainingArgs);
+            if (logs === undefined) {
+                throw new Error("Epoch didn't gave any logs");
+            }
+            const { loss, val_acc, val_loss, peakMemory } = logs;
+            if (loss === undefined || isNaN(loss)) {
+                throw new Error("Training loss is undefined or nan");
+            }
+            const structuredLogs = {
+                epoch,
+                peakMemory,
+                training: {
+                    loss: logs.loss
+                }
+            };
+            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);
+    }
+    async generate(input, tokenizer, newTokens = 10) {
+        const { input_ids: tokens } = await tokenizer(input, { return_tensor: false });
+        const generationConfig = {
+            maxNewTokens: newTokens,
+            temperature: 1.0,
+            doSample: false
+        };
+        const predictedTokens = await this.model.generate(tokens, generationConfig);
+        const generatedWords = tokenizer.decode(predictedTokens[0]);
+        return generatedWords;
+    }
+    get config() {
+        return this.model.getGPTConfig;
+    }
+    get weights() {
+        return new WeightsContainer(this.model.weights.map((w) => w.read()));
+    }
+    set weights(ws) {
+        this.model.setWeights(ws.weights);
+    }
+    static deserialize(data) {
+        const model = new GPT(data.config);
+        model.weights = data.weights;
+        return model;
+    }
+    serialize() {
+        return {
+            weights: this.weights,
+            config: this.config
+        };
+    }
+    [Symbol.dispose]() {
+        console.log("Disposing model");
+        if (this.model.optimizer !== undefined) {
+            this.model.optimizer.dispose();
+        }
+        // Some tensors are not cleaned up when model.dispose is called 
+        // So we dispose them manually
+        this.model.disposeRefs();
+        this.model.dispose();
+    }
+}

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

@@ -0,0 +1,13 @@
+import * as tf from '@tensorflow/tfjs';
+import type { GPTConfig } from './config.js';
+/**
+ * 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 declare function GPTArchitecture(config: Required<GPTConfig>, disposalRefs: tf.TensorContainer[], peakMemory: {
+    value: number;
+}): tf.LayersModel;

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