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

package/dist/dataset/dataset_builder.d.ts

@@ -17,15 +17,11 @@ export declare class DatasetBuilder<Source> {
     /**
      * The buffer of unlabelled file sources.
      */
-    private _sources;
+    private _unlabeledSources;
     /**
      * The buffer of labelled file sources.
      */
-    private labelledSources;
-    /**
-     * Whether a dataset was already produced.
-     */
-    private _built;
+    private _labeledSources;
     constructor(
     /**
      * The data loader used to load the data contained in the provided files.
@@ -48,13 +44,8 @@ export declare class DatasetBuilder<Source> {
      * @param label The file sources label
      */
     clearFiles(label?: string): void;
-    private resetBuiltState;
     private getLabels;
     build(config?: DataConfig): Promise<DataSplit>;
-    /**
-     * Whether the dataset builder has already been consumed to produce a dataset.
-     */
-    get built(): boolean;
     get size(): number;
     get sources(): Source[];
 }

package/dist/dataset/dataset_builder.js

@@ -9,16 +9,11 @@ export class DatasetBuilder {
     /**
      * The buffer of unlabelled file sources.
      */
-    _sources;
+    _unlabeledSources;
     /**
      * The buffer of labelled file sources.
      */
-    labelledSources;
-    /**
-     * Whether a dataset was already produced.
-     */
-    // TODO useless, responsibility on callers
-    _built;
+    _labeledSources;
     constructor(
     /**
      * The data loader used to load the data contained in the provided files.
@@ -30,9 +25,9 @@ export class DatasetBuilder {
     task) {
         this.dataLoader = dataLoader;
         this.task = task;
-        this._sources = [];
-        this.labelledSources = Map();
-        this._built = false;
+        this._unlabeledSources = [];
+        // Map from label to sources
+        this._labeledSources = Map();
     }
     /**
      * Adds the given file sources to the builder's buffer. Sources may be provided a label in the case
@@ -41,19 +36,16 @@ export class DatasetBuilder {
      * @param label The file sources label
      */
     addFiles(sources, label) {
-        if (this.built) {
-            this.resetBuiltState();
-        }
         if (label === undefined) {
-            this._sources = this._sources.concat(sources);
+            this._unlabeledSources = this._unlabeledSources.concat(sources);
         }
         else {
-            const currentSources = this.labelledSources.get(label);
+            const currentSources = this._labeledSources.get(label);
             if (currentSources === undefined) {
-                this.labelledSources = this.labelledSources.set(label, sources);
+                this._labeledSources = this._labeledSources.set(label, sources);
             }
             else {
-                this.labelledSources = this.labelledSources.set(label, currentSources.concat(sources));
+                this._labeledSources = this._labeledSources.set(label, currentSources.concat(sources));
             }
         }
     }
@@ -63,27 +55,19 @@ export class DatasetBuilder {
      * @param label The file sources label
      */
     clearFiles(label) {
-        if (this.built) {
-            this.resetBuiltState();
-        }
         if (label === undefined) {
-            this._sources = [];
+            this._unlabeledSources = [];
         }
         else {
-            this.labelledSources = this.labelledSources.delete(label);
+            this._labeledSources = this._labeledSources.delete(label);
         }
     }
-    // If files are added or removed, then this should be called since the latest
-    // version of the dataset_builder has not yet been built.
-    resetBuiltState() {
-        this._built = false;
-    }
     getLabels() {
         // We need to duplicate the labels as we need one for each source.
         // Say for label A we have sources [img1, img2, img3], then we
         // need labels [A, A, A].
         let labels = [];
-        this.labelledSources.forEach((sources, label) => {
+        this._labeledSources.forEach((sources, label) => {
             const sourcesLabels = Array.from({ length: sources.length }, (_) => label);
             labels = labels.concat(sourcesLabels);
         });
@@ -91,17 +75,17 @@ export class DatasetBuilder {
     }
     async build(config) {
         // Require that at least one source collection is non-empty, but not both
-        if ((this._sources.length > 0) === (this.labelledSources.size > 0)) {
-            throw new Error('Please provide dataset input files'); // This error message is parsed in DatasetInput.vue
+        if (this._unlabeledSources.length + this._labeledSources.size === 0) {
+            throw new Error('No input files connected'); // This error message is parsed in Trainer.vue
         }
         let dataTuple;
-        if (this._sources.length > 0) {
+        if (this._unlabeledSources.length > 0) {
             let defaultConfig = {};
             if (config?.inference === true) {
                 // Inferring model, no labels needed
                 defaultConfig = {
                     features: this.task.trainingInformation.inputColumns,
-                    shuffle: false
+                    shuffle: true
                 };
             }
             else {
@@ -109,34 +93,26 @@ export class DatasetBuilder {
                 defaultConfig = {
                     features: this.task.trainingInformation.inputColumns,
                     labels: this.task.trainingInformation.outputColumns,
-                    shuffle: false
+                    shuffle: true
                 };
             }
-            dataTuple = await this.dataLoader.loadAll(this._sources, { ...defaultConfig, ...config });
+            dataTuple = await this.dataLoader.loadAll(this._unlabeledSources, { ...defaultConfig, ...config });
         }
         else {
             // Labels are inferred from the file selection boxes
             const defaultConfig = {
                 labels: this.getLabels(),
-                shuffle: false
+                shuffle: true
             };
-            const sources = this.labelledSources.valueSeq().toArray().flat();
+            const sources = this._labeledSources.valueSeq().toArray().flat();
             dataTuple = await this.dataLoader.loadAll(sources, { ...defaultConfig, ...config });
         }
-        // TODO @s314cy: Support .csv labels for image datasets (supervised training or testing)
-        this._built = true;
         return dataTuple;
     }
-    /**
-     * Whether the dataset builder has already been consumed to produce a dataset.
-     */
-    get built() {
-        return this._built;
-    }
     get size() {
-        return Math.max(this._sources.length, this.labelledSources.size);
+        return Math.max(this._unlabeledSources.length, this._labeledSources.size);
     }
     get sources() {
-        return this._sources.length > 0 ? this._sources : this.labelledSources.valueSeq().toArray().flat();
+        return this._unlabeledSources.length > 0 ? this._unlabeledSources : this._labeledSources.valueSeq().toArray().flat();
     }
 }

package/dist/default_tasks/cifar10.d.ts

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

package/dist/default_tasks/{cifar10/index.js → cifar10.js}

@@ -1,6 +1,6 @@
 import * as tf from '@tensorflow/tfjs';
-import { data, models } from '../../index.js';
-import baseModel from './model.js';
+import { data, models } from '../index.js';
+import baseModel from '../models/mobileNet_v1_025_224.js';
 export const cifar10 = {
     getTask() {
         return {

package/dist/default_tasks/index.d.ts

@@ -1,6 +1,7 @@
-export { cifar10 } from './cifar10/index.js';
+export { cifar10 } from './cifar10.js';
 export { lusCovid } from './lus_covid.js';
+export { skinCondition } from './skin_condition.js';
 export { mnist } from './mnist.js';
-export { simpleFace } from './simple_face/index.js';
+export { simpleFace } from './simple_face.js';
 export { titanic } from './titanic.js';
 export { wikitext } from './wikitext.js';

package/dist/default_tasks/index.js

@@ -1,6 +1,7 @@
-export { cifar10 } from './cifar10/index.js';
+export { cifar10 } from './cifar10.js';
 export { lusCovid } from './lus_covid.js';
+export { skinCondition } from './skin_condition.js';
 export { mnist } from './mnist.js';
-export { simpleFace } from './simple_face/index.js';
+export { simpleFace } from './simple_face.js';
 export { titanic } from './titanic.js';
 export { wikitext } from './wikitext.js';

package/dist/default_tasks/lus_covid.js

@@ -8,7 +8,7 @@ export const lusCovid = {
                 taskTitle: 'COVID Lung Ultrasound',
                 summary: {
                     preview: 'Do you have a data of lung ultrasound images on patients <b>suspected of Lower Respiratory Tract infection (LRTI) during the COVID pandemic</b>? <br> Learn how to discriminate between COVID positive and negative patients by joining this task.',
-                    overview: "Don’t have a dataset of your own? Download a sample of a few cases <a class='underline' href='https://drive.switch.ch/index.php/s/zM5ZrUWK3taaIly' target='_blank'>here</a>."
+                    overview: "Don't have a dataset of your own? Download a sample of a few cases <a class='underline' href='https://drive.switch.ch/index.php/s/zM5ZrUWK3taaIly' target='_blank'>here</a>."
                 },
                 model: "We use a simplified* version of the <b>DeepChest model</b>: A deep learning model developed in our lab (<a class='underline' href='https://www.epfl.ch/labs/mlo/igh-intelligent-global-health/'>intelligent Global Health</a>.). On a cohort of 400 Swiss patients suspected of LRTI, the model obtained over 90% area under the ROC curve for this task. <br><br>*Simplified to ensure smooth running on your browser, the performance is minimally affected. Details of the adaptations are below <br>- <b>Removed</b>: positional embedding (i.e. we don’t take the anatomic position into consideration). Rather, the model now does mean pooling over the feature vector of the images for each patient <br>- <b>Replaced</b>: ResNet18 by Mobilenet",
                 dataFormatInformation: 'This model takes as input an image dataset. It consists on a set of lung ultrasound images per patient with its corresponding label of covid positive or negative. Moreover, to identify the images per patient you have to follow the follwing naming pattern: "patientId_*.png"',

package/dist/default_tasks/simple_face.d.ts

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

package/dist/default_tasks/{simple_face/index.js → simple_face.js}

@@ -1,6 +1,6 @@
 import * as tf from '@tensorflow/tfjs';
-import { data, models } from '../../index.js';
-import baseModel from './model.js';
+import { data, models } from '../index.js';
+import baseModel from '../models/mobileNetV2_35_alpha_2_classes.js';
 export const simpleFace = {
     getTask() {
         return {
@@ -12,7 +12,7 @@ export const simpleFace = {
                     overview: 'Simple face is a small subset of face_task from Kaggle'
                 },
                 dataFormatInformation: '',
-                dataExampleText: 'Below you find an example',
+                dataExampleText: 'Below you can find an example',
                 dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/simple_face-example.png'
             },
             trainingInformation: {

package/dist/default_tasks/skin_condition.d.ts

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

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;