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

package/README.md

@@ -0,0 +1,24 @@
+# @epfml/discojs
+
+Decentralized & federated privacy-preserving ML training in TypeScript.
+
+This is the core library of the Disco.js project.
+
+It is platform-agnostic, and has two companions library:
+ - [`discojs-node`](../discojs-node) for Node.js
+ - [`discojs-web`](../discojs-web) for web browsers
+
+The easiest way to start using it is through the `Disco` object.
+Create your own `Task` or load one from our `default_tasks`,
+setup the `Dataset` you want, and train with it.
+
+```ts
+import { Disco } from '@epfml/discojs'
+
+const url = ...; // url to a Disco.js server
+const dataset = ...;
+const task = ...;
+
+const disco = new Disco(task, { url })
+for await (const _ of disco.fit(dataset));
+```

package/dist/dataset/dataset_builder.js

@@ -92,7 +92,7 @@ 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');
+            throw new Error('Please provide dataset input files'); // This error message is parsed in DatasetInput.vue
         }
         let dataTuple;
         if (this._sources.length > 0) {

package/dist/default_tasks/cifar10/index.js

@@ -11,11 +11,10 @@ export const cifar10 = {
                     preview: 'In this challenge, we ask you to classify images into categories based on the objects shown on the image.',
                     overview: 'The CIFAR-10 dataset is a collection of images that are commonly used to train machine learning and computer vision algorithms. It is one of the most widely used datasets for machine learning research.'
                 },
-                limitations: 'The training data is limited to small images of size 32x32.',
-                tradeoffs: 'Training success strongly depends on label distribution',
                 dataFormatInformation: 'Images should be of .png format and of size 32x32. <br> The label file should be .csv, where each row contains a file_name, class.  <br> <br> e.g. if you have images: 0.png (of a frog) and 1.png (of a car) <br> labels.csv contains: (Note that no header is needed)<br> 0.png, frog <br> 1.png, car',
                 dataExampleText: 'Below you can find 10 random examples from each of the 10 classes in the dataset.',
-                dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/cifar10-example.png'
+                dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/cifar10-example.png',
+                sampleDatasetLink: 'https://www.kaggle.com/competitions/cifar-10/data'
             },
             trainingInformation: {
                 modelID: 'cifar10-model',
@@ -27,7 +26,7 @@ export const cifar10 = {
                 preprocessingFunctions: [data.ImagePreprocessing.Resize],
                 IMAGE_H: 224,
                 IMAGE_W: 224,
-                LABEL_LIST: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'],
+                LABEL_LIST: ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'],
                 scheme: 'decentralized',
                 noiseScale: undefined,
                 clippingRadius: 20,

package/dist/default_tasks/index.d.ts

@@ -1,8 +1,6 @@
 export { cifar10 } from './cifar10/index.js';
-export { geotags } from './geotags/index.js';
 export { lusCovid } from './lus_covid.js';
 export { mnist } from './mnist.js';
 export { simpleFace } from './simple_face/index.js';
-export { skinMnist } from './skin_mnist.js';
 export { titanic } from './titanic.js';
 export { wikitext } from './wikitext.js';

package/dist/default_tasks/index.js

@@ -1,8 +1,6 @@
 export { cifar10 } from './cifar10/index.js';
-export { geotags } from './geotags/index.js';
 export { lusCovid } from './lus_covid.js';
 export { mnist } from './mnist.js';
 export { simpleFace } from './simple_face/index.js';
-export { skinMnist } from './skin_mnist.js';
 export { titanic } from './titanic.js';
 export { wikitext } from './wikitext.js';

package/dist/default_tasks/lus_covid.js

@@ -11,16 +11,16 @@ export const lusCovid = {
                     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",
-                tradeoffs: 'We are using a simpler version of DeepChest in order to be able to run it on the browser.',
                 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"',
                 dataExampleText: 'Below you can find an example of an expected lung image for patient 2 named: 2_QAID_1.masked.reshaped.squared.224.png',
-                dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/2_QAID_1.masked.reshaped.squared.224.png'
+                dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/2_QAID_1.masked.reshaped.squared.224.png',
+                sampleDatasetLink: 'https://drive.switch.ch/index.php/s/zM5ZrUWK3taaIly'
             },
             trainingInformation: {
                 modelID: 'lus-covid-model',
                 epochs: 50,
                 roundDuration: 2,
-                validationSplit: 0,
+                validationSplit: 0.2,
                 batchSize: 5,
                 IMAGE_H: 100,
                 IMAGE_W: 100,

package/dist/default_tasks/mnist.js

@@ -11,7 +11,6 @@ export const mnist = {
                     overview: 'The MNIST handwritten digit classification problem is a standard dataset used in computer vision and deep learning. Although the dataset is effectively solved, we use it to test our Decentralised Learning algorithms and platform.'
                 },
                 model: 'The current model is a very simple CNN and its main goal is to test the app and the Decentralizsed Learning functionality.',
-                tradeoffs: 'We are using a simple model, first a 2d convolutional layer > max pooling > 2d convolutional layer > max pooling > convolutional layer > 2 dense layers.',
                 dataFormatInformation: 'This model is trained on images corresponding to digits 0 to 9. You can upload each digit image of your dataset in the box corresponding to its label. The model taskes images of size 28x28 as input.',
                 dataExampleText: 'Below you can find an example of an expected image representing the digit 9.',
                 dataExampleImage: 'http://storage.googleapis.com/deai-313515.appspot.com/example_training_data/9-mnist-example.png'

package/dist/default_tasks/simple_face/index.js

@@ -11,8 +11,6 @@ export const simpleFace = {
                     preview: 'Can you detect if the person in a picture is a child or an adult?',
                     overview: 'Simple face is a small subset of face_task from Kaggle'
                 },
-                limitations: 'The training data is limited to small images of size 200x200.',
-                tradeoffs: 'Training success strongly depends on label distribution',
                 dataFormatInformation: '',
                 dataExampleText: 'Below you find an example',
                 dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/simple_face-example.png'

package/dist/default_tasks/titanic.js

@@ -11,7 +11,6 @@ export const titanic = {
                     overview: 'We all know the unfortunate story of the Titanic: this flamboyant new transatlantic boat that sunk in 1912 in the North Atlantic Ocean. Today, we revist this tragedy by trying to predict the survival odds of the passenger given some basic features.'
                 },
                 model: 'The current model does not normalize the given data and applies only a very simple pre-processing of the data.',
-                tradeoffs: 'We are using a small model for this task: 4 fully connected layers with few neurons. This allows fast training but can yield to reduced accuracy.',
                 dataFormatInformation: 'This model takes as input a CSV file with 12 columns. The features are general information about the passenger (sex, age, name, etc.) and specific related Titanic data such as the ticket class bought by the passenger, its cabin number, etc.<br><br>pclass: A proxy for socio-economic status (SES)<br>1st = Upper<br>2nd = Middle<br>3rd = Lower<br><br>age: Age is fractional if less than 1. If the age is estimated, it is in the form of xx.5<br><br>sibsp: The dataset defines family relations in this way:<br>Sibling = brother, sister, stepbrother, stepsister<br>Spouse = husband, wife (mistresses and fiancés were ignored)<br><br>parch: The dataset defines family relations in this way:<br>Parent = mother, father<br>Child = daughter, son, stepdaughter, stepson<br>Some children travelled only with a nanny, therefore parch=0 for them.<br><br>The first line of the CSV contains the header:<br> PassengerId, Survived, Pclass, Name, Sex, Age, SibSp, Parch, Ticket, Fare, Cabin, Embarked<br><br>Each susequent row contains the corresponding data.',
                 dataExampleText: 'Below one can find an example of a datapoint taken as input by our model. In this datapoint, the person is young man named Owen Harris that unfortunnalty perished with the Titanic. He boarded the boat in South Hamptons and was a 3rd class passenger. On the testing & validation page, the data should not contain the label column (Survived).',
                 dataExample: [

package/dist/default_tasks/wikitext.js

@@ -9,10 +9,9 @@ export const wikitext = {
                     preview: 'In this challenge, we ask you to do next word prediction on a dataset of Wikipedia articles.',
                     overview: 'Wikitext-103-raw is a dataset comprising unprocessed text excerpts from Wikipedia articles, designed for tasks related to natural language processing and language modeling.'
                 },
-                limitations: 'The dataset may contain noise, inconsistencies, and unstructured content due to its raw nature, potentially posing challenges for certain NLP tasks.',
-                tradeoffs: 'The raw format may lack structured annotations and may require additional preprocessing for specific applications.',
                 dataFormatInformation: 'The dataset is organized as a large text file, with each line representing a segment of raw text from Wikipedia articles.',
-                dataExampleText: 'An example excerpt from the dataset could be: "The history of artificial intelligence dates back to ancient times, with philosophical discussions on the nature of thought and reasoning."'
+                dataExampleText: 'An example excerpt from the dataset could be: "The history of artificial intelligence dates back to ancient times, with philosophical discussions on the nature of thought and reasoning."',
+                sampleDatasetLink: 'https://dax-cdn.cdn.appdomain.cloud/dax-wikitext-103/1.0.1/wikitext-103.tar.gz'
             },
             trainingInformation: {
                 dataType: 'text',

package/dist/index.d.ts

@@ -8,7 +8,7 @@ export * as aggregator from './aggregator/index.js';
 export { WeightsContainer, aggregation } from './weights/index.js';
 export { AsyncInformant } from './async_informant.js';
 export { Logger, ConsoleLogger } from './logging/index.js';
-export { Memory, ModelType, type ModelInfo, type Path, type ModelSource, Empty as EmptyMemory } from './memory/index.js';
+export { Memory, StoredModelType, type ModelInfo, type Path, type ModelSource, Empty as EmptyMemory } from './memory/index.js';
 export { Disco, RoundLogs } from './training/index.js';
 export { Validator } from './validation/index.js';
 export { Model, EpochLogs } from './models/index.js';

package/dist/index.js

@@ -8,7 +8,7 @@ export * as aggregator from './aggregator/index.js';
 export { WeightsContainer, aggregation } from './weights/index.js';
 export { AsyncInformant } from './async_informant.js';
 export { ConsoleLogger } from './logging/index.js';
-export { Memory, ModelType, Empty as EmptyMemory } from './memory/index.js';
+export { Memory, StoredModelType, Empty as EmptyMemory } from './memory/index.js';
 export { Disco } from './training/index.js';
 export { Validator } from './validation/index.js';
 export { Model } from './models/index.js';

package/dist/memory/base.d.ts

@@ -1,5 +1,5 @@
 import type { Model, TaskID } from '../index.js';
-import type { ModelType } from './model_type.js';
+import type { StoredModelType } from './model_type.js';
 /**
  * Model path which uniquely identifies a model in memory.
  */
@@ -8,22 +8,9 @@ export type Path = string;
  * Model information which uniquely identifies a model in memory.
  */
 export interface ModelInfo {
-    /**
-     * The model's type: "working" or "saved" model.
-     */
-    type?: ModelType;
-    /**
-     * The model's version, to allow for multiple saved models of a same task without
-     * causing id conflicts
-     */
+    type?: StoredModelType;
     version?: number;
-    /**
-     * The model's corresponding task
-     */
     taskID: TaskID;
-    /**
-     * The model's name
-     */
     name: string;
 }
 /**
@@ -95,21 +82,21 @@ export declare abstract class Memory {
     /**
      * Computes the path in memory corresponding to the given model source, be it a path or model information.
      * This is used to easily switch between model path and information, which are both unique model identifiers
-     * with a one-to-one correspondance. Returns undefined instead if no path could be inferred from the given
+     * with a one-to-one equivalence. Returns undefined instead if no path could be inferred from the given
      * model source.
      * @param source The model source
      * @returns The model path
      */
-    abstract pathFor(source: ModelSource): Path | undefined;
+    abstract getModelMemoryPath(source: ModelSource): Path | undefined;
     /**
      * Computes the model information corresponding to the given model source, be it a path or model information.
      * This is used to easily switch between model path and information, which are both unique model identifiers
-     * with a one-to-one correspondance. Returns undefined instead if no unique model information could be inferred
+     * with a one-to-one equivalence. Returns undefined instead if no unique model information could be inferred
      * from the given model source.
      * @param source The model source
      * @returns The model information
      */
-    abstract infoFor(source: ModelSource): ModelInfo | undefined;
+    abstract getModelInfo(source: ModelSource): ModelInfo | undefined;
     /**
      * Computes the lowest version a model source can have without conflicting with model versions currently in memory.
      * @param source The model source

package/dist/memory/empty.d.ts

@@ -14,7 +14,7 @@ export declare class Empty extends Memory {
     saveModel(): Promise<undefined>;
     deleteModel(): Promise<void>;
     downloadModel(): Promise<void>;
-    pathFor(): Path;
-    infoFor(): ModelInfo;
+    getModelMemoryPath(): Path;
+    getModelInfo(): ModelInfo;
     duplicateSource(): Promise<undefined>;
 }

package/dist/memory/empty.js

@@ -31,10 +31,10 @@ export class Empty extends Memory {
     downloadModel() {
         return Promise.reject(new Error('empty'));
     }
-    pathFor() {
+    getModelMemoryPath() {
         throw new Error('empty');
     }
-    infoFor() {
+    getModelInfo() {
         throw new Error('empty');
     }
     duplicateSource() {

package/dist/memory/index.d.ts

@@ -1,3 +1,3 @@
 export { Empty } from './empty.js';
 export { Memory, type ModelInfo, type Path, type ModelSource } from './base.js';
-export { ModelType } from './model_type.js';
+export { StoredModelType } from './model_type.js';

package/dist/memory/index.js

@@ -1,3 +1,3 @@
 export { Empty } from './empty.js';
 export { Memory } from './base.js';
-export { ModelType } from './model_type.js';
+export { StoredModelType } from './model_type.js';

package/dist/memory/model_type.d.ts

@@ -3,7 +3,7 @@
  * being trained ("working model") or a regular model saved in memory ("saved model").
  * There can only be a single working model for a given task.
  */
-export declare enum ModelType {
+export declare enum StoredModelType {
     WORKING = "working",
     SAVED = "saved"
 }

package/dist/memory/model_type.js

@@ -3,8 +3,8 @@
  * being trained ("working model") or a regular model saved in memory ("saved model").
  * There can only be a single working model for a given task.
  */
-export var ModelType;
-(function (ModelType) {
-    ModelType["WORKING"] = "working";
-    ModelType["SAVED"] = "saved";
-})(ModelType || (ModelType = {}));
+export var StoredModelType;
+(function (StoredModelType) {
+    StoredModelType["WORKING"] = "working";
+    StoredModelType["SAVED"] = "saved";
+})(StoredModelType || (StoredModelType = {}));

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

@@ -0,0 +1,32 @@
+type GPTModelType = 'gpt2' | 'gpt2-medium' | 'gpt2-large' | 'gpt2-xl' | 'gpt-mini' | 'gpt-micro' | 'gpt-nano';
+export interface GPTConfig {
+    lr: number;
+    blockSize: number;
+    vocabSize: number;
+    modelType: GPTModelType;
+    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: GPTModelType): 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,37 @@
+/**
+ * this code is taken from gpt-tfjs with modifications from @peacefulotter and @lukemovement
+ **/
+import * as tf from '@tensorflow/tfjs';
+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;
+    extract(): tf.LayersModel;
+    [Symbol.dispose](): void;
+}
+export type GPTSerialization = {
+    weights: WeightsContainer;
+    config?: GPTConfig;
+};

package/dist/models/gpt/index.js

@@ -0,0 +1,107 @@
+/**
+ * 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
+        };
+    }
+    extract() {
+        return this.model;
+    }
+    [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;