@epfml/discojs 2.1.2-p20240624145915.0 → 2.1.2-p20240702170238.0

package/dist/client/decentralized/peer.d.ts

@@ -1,4 +1,3 @@
-/// <reference types="node" resolution-mode="require"/>
 import type { NodeID } from '../types.js';
 export type SignalData = {
     type: 'answer' | 'offer' | 'pranswer' | 'rollback';

package/dist/default_tasks/cifar10.js

@@ -8,13 +8,15 @@ export const cifar10 = {
             displayInformation: {
                 taskTitle: 'CIFAR10',
                 summary: {
-                    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.'
+                    preview: 'CIFAR-10 is a classic image classification task, and one of the most widely used datasets for machine learning research.',
+                    overview: "The dataset contains 60,000 32x32 color images in 10 different classes: airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. The official CIFAR-10 website can be found <a class='underline text-blue-400' href='https://www.cs.toronto.edu/~kriz/cifar.html' target='_blank'>here</a>. You can find a link to a sample dataset at the next step (Connect Your Data)."
                 },
-                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',
+                model: 'The model is a pretrained <a  class="underline text-blue-400" target="_blank" href="https://github.com/tensorflow/tfjs-models/tree/master/mobilenet">MobileNetV1 model</a> trained in Tensorflow.js. The last output layer is replaced with a fully connected layer with softmax activation and one output neuron per CIFAR10 category. The data preprocessing reshapes images into 224x224 pixels and normalizes values between 0 and 1. The neural network is optimized via Stochastic Gradient Descent and a categorical Cross Entropy loss.',
+                dataFormatInformation: 'Images should be of .png format and of size 32x32. <br> The CSV file should start with the exact header "filename,label", and each row should contain an image filename (without extension) and its label.<br><br> For example if you have images: 0.png (of a frog) and 1.png (of a car) <br> The CSV file should be: <br>filename, label <br><br> 0, frog <br> 1, 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',
-                sampleDatasetLink: 'https://www.kaggle.com/competitions/cifar-10/data'
+                sampleDatasetLink: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data.tar.gz',
+                sampleDatasetInstructions: 'Opening the link should start downloading a zip file which you can unzip. To connect the data, use the CSV option below and select the file named "cifar10-labels.csv". You can now connect the images located in the "CIFAR10" folder. Note that there are only 24 images in this sample dataset which is far too few to successfully train a machine learning model.'
             },
             trainingInformation: {
                 modelID: 'cifar10-model',
@@ -23,7 +25,7 @@ export const cifar10 = {
                 validationSplit: 0.2,
                 batchSize: 10,
                 dataType: 'image',
-                preprocessingFunctions: [data.ImagePreprocessing.Resize],
+                preprocessingFunctions: [data.ImagePreprocessing.Resize, data.ImagePreprocessing.Normalize],
                 IMAGE_H: 224,
                 IMAGE_W: 224,
                 LABEL_LIST: ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'],

package/dist/default_tasks/index.d.ts

@@ -1,6 +1,5 @@
 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.js';
 export { titanic } from './titanic.js';

package/dist/default_tasks/index.js

@@ -1,6 +1,5 @@
 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.js';
 export { titanic } from './titanic.js';

package/dist/default_tasks/lus_covid.js

@@ -5,16 +5,17 @@ export const lusCovid = {
         return {
             id: 'lus_covid',
             displayInformation: {
-                taskTitle: 'COVID Lung Ultrasound',
+                taskTitle: 'COVID-19 Diagnosis from Lung Ultrasounds',
                 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>."
+                    preview: "Medical images are a typical example of data that exists in huge quantity yet that can't be shared due to confidentiality reasons. Medical applications would immensely benefit from training on data currently locked. More data diversity leads to better generalization and bias mitigation.",
+                    overview: "Disco allows data owners to collaboratively train machine learning models using their respective data without any privacy breach. This example problem is about diagnosing whether patients are positive or negative to COVID-19 from lung ultrasounds images. <br>Don't have a dataset of your own? You can find a link to a sample dataset at the next step."
                 },
-                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"',
-                dataExampleText: 'Below you can find an example of an expected lung image for patient 2 named: 2_QAID_1.masked.reshaped.squared.224.png',
+                model: "The model is a simple Convolutional Neural Network composed of two convolutional layers with ReLU activations and max pooling layers, followed by a fully connected output layer. The data preprocessing reshapes images into 100x100 pixels and normalizes values between 0 and 1",
+                dataFormatInformation: 'This model takes as input an image dataset of lung ultrasounds. The images are resized automatically.',
+                dataExampleText: 'Below you can find an example of an expected lung image.',
                 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'
+                sampleDatasetLink: 'https://drive.switch.ch/index.php/s/zM5ZrUWK3taaIly',
+                sampleDatasetInstructions: 'Opening the link will take you to a Switch Drive folder. You can click on the Download button in the top right corner. Unzip the file and you will get two subfolders: "COVID-" and "COVID+". You can connect the data by using the Group option and selecting each image group in its respective field.'
             },
             trainingInformation: {
                 modelID: 'lus-covid-model',

package/dist/default_tasks/mnist.js

@@ -1,30 +1,32 @@
 import * as tf from '@tensorflow/tfjs';
-import { models } from '../index.js';
+import { data, models } from '../index.js';
 export const mnist = {
     getTask() {
         return {
             id: 'mnist',
             displayInformation: {
-                taskTitle: 'MNIST',
+                taskTitle: 'Handwritten Digit Recognition',
                 summary: {
-                    preview: "Test our platform by using a publicly available <b>image</b> dataset. <br><br> Download the classic MNIST imagebank of hand-written numbers <a class='underline text-primary-dark dark:text-primary-light' href='https://www.kaggle.com/scolianni/mnistasjpg'>here</a>. <br> This model learns to identify hand written numbers.",
-                    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.'
+                    preview: "The MNIST handwritten digit classification problem is a classic dataset used in computer vision and deep learning. The objective is to classify handwritten digits from 28x28 pixel images.",
+                    overview: "Download the classic MNIST dataset of hand-written numbers <a class='underline text-blue-400' target='_blank' href='https://www.kaggle.com/scolianni/mnistasjpg'>here</a>. You can also find a sample dataset at the next step."
                 },
-                model: 'The current model is a very simple CNN and its main goal is to test the app and the Decentralizsed Learning functionality.',
-                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.',
+                model: "The model is a simple Convolutional Neural Network composed of three convolutional layers with ReLU activations and max pooling layers, followed by two fully connected layers. The data preprocessing simply normalizes values between 0 and 1. The neural network is optimized via RMSProp and a categorical cross-entropy loss.",
+                dataFormatInformation: 'This model is trained on images corresponding to digits 0 to 9. You can connect your own images of each digit in the box corresponding to its label. The model takes 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'
+                dataExampleImage: 'http://storage.googleapis.com/deai-313515.appspot.com/example_training_data/9-mnist-example.png',
+                sampleDatasetLink: 'https://storage.googleapis.com/deai-313515.appspot.com/MNIST_samples.tar.gz',
+                sampleDatasetInstructions: 'Opening the link should start downloading a zip file which you can unzip. You can connect the data with the CSV option below using the CSV file named "mnist_labels.csv". After selecting in the CSV file, you will be able to connect the data under in the "images" folder.'
             },
             trainingInformation: {
                 modelID: 'mnist-model',
-                epochs: 10,
+                epochs: 20,
                 roundDuration: 10,
                 validationSplit: 0.2,
                 batchSize: 30,
                 dataType: 'image',
                 IMAGE_H: 28,
                 IMAGE_W: 28,
-                preprocessingFunctions: [],
+                preprocessingFunctions: [data.ImagePreprocessing.Normalize],
                 LABEL_LIST: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'],
                 scheme: 'decentralized',
                 noiseScale: undefined,

package/dist/default_tasks/simple_face.js

@@ -9,11 +9,13 @@ export const simpleFace = {
                 taskTitle: 'Simple Face',
                 summary: {
                     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'
+                    overview: 'Simple face is a small subset of the public face_task dataset from Kaggle'
                 },
                 dataFormatInformation: '',
                 dataExampleText: 'Below you can find an example',
-                dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/simple_face-example.png'
+                dataExampleImage: 'https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/simple_face-example.png',
+                sampleDatasetLink: "https://storage.googleapis.com/deai-313515.appspot.com/example_training_data.tar.gz",
+                sampleDatasetInstructions: 'Opening the link should start downloading a zip file which you can unzip. Inside the "example_training_data" directory you should find the "simple_face" folder which contains the "adult" and "child" folders. To connect the data, select the Group option below and connect adults and children image groups.'
             },
             trainingInformation: {
                 modelID: 'simple_face-model',

package/dist/default_tasks/titanic.js

@@ -5,14 +5,14 @@ export const titanic = {
         return {
             id: 'titanic',
             displayInformation: {
-                taskTitle: 'Titanic',
+                taskTitle: 'Titanic Prediction',
                 summary: {
-                    preview: "Test our platform by using a publicly available <b>tabular</b> dataset. <br><br> Download the passenger list from the Titanic shipwreck here: <a class='underline text-primary-dark dark:text-primary-light' href='https://github.com/epfml/disco/raw/develop/example_training_data/titanic_train.csv'>titanic_train.csv</a> (more info <a class='underline text-primary-dark dark:text-primary-light' href='https://www.kaggle.com/c/titanic'>here</a>). <br> This model predicts the type of person most likely to survive/die in the historic ship accident, based on their characteristics (sex, age, class etc.).",
-                    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.'
+                    preview: "The Titanic classification task is one of the main entrypoints into machine learning. Using passenger data (name, age, gender, socio-economic class, etc), the goal is to identify who was more likely to survive the infamous shipwreck.",
+                    overview: "The original competition can be found on <a  target='_blank' class='underline text-blue-400' href='https://www.kaggle.com/c/titanic'>Kaggle</a> and a link to the training set can be found here <a target='_blank' class='underline text-blue-400' href='https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/titanic_train.csv'>here</a>."
                 },
-                model: 'The current model does not normalize the given data and applies only a very simple pre-processing of the data.',
-                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).',
+                model: 'The model is a simple 5-layer feedforward network with ReLU activations. The model is optimized with Adam and binary cross-entropy loss. The preprocessing only fills missing value with a placeholder value (0).',
+                dataFormatInformation: 'The expected format for the tabular dataset is exactly the same as the sample data provided above or in the Kaggle competition. It is 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>The first line of the CSV contains the header: "PassengerId, Survived, Pclass, Name, Sex, Age, SibSp, Parch, Ticket, Fare, Cabin, Embarked"<br>Each subsequent row contains passenger data.',
+                dataExampleText: "Here's an example of one data point:",
                 dataExample: [
                     { columnName: 'PassengerId', columnData: '1' },
                     { columnName: 'Survived', columnData: '0' },
@@ -40,11 +40,13 @@ export const titanic = {
                     'Cabin',
                     'Embarked',
                     'Pclass'
-                ]
+                ],
+                sampleDatasetLink: "https://storage.googleapis.com/deai-313515.appspot.com/example_training_data/titanic_train.csv",
+                sampleDatasetInstructions: 'Opening the link should start downloading a CSV file which you can drag and drop in the field below.'
             },
             trainingInformation: {
                 modelID: 'titanic-model',
-                epochs: 20,
+                epochs: 40,
                 roundDuration: 10,
                 validationSplit: 0.2,
                 batchSize: 30,
@@ -79,7 +81,7 @@ export const titanic = {
         model.add(tf.layers.dense({ units: 32, activation: 'relu' }));
         model.add(tf.layers.dense({ units: 1, activation: 'sigmoid' }));
         model.compile({
-            optimizer: tf.train.sgd(0.001),
+            optimizer: 'adam',
             loss: 'binaryCrossentropy',
             metrics: ['accuracy']
         });

package/dist/default_tasks/wikitext.js

@@ -2,28 +2,30 @@ import { data, models } from '../index.js';
 export const wikitext = {
     getTask() {
         return {
-            id: 'wikitext-103',
+            id: 'llm_task',
             displayInformation: {
-                taskTitle: 'Language modelling on wikitext',
+                taskTitle: 'GPT Language Modeling',
                 summary: {
-                    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.'
+                    preview: 'Train a language model (L)LM in your browser, collaboratively and from scratch.',
+                    overview: "You can train a GPT-2 model in your browser and in a collaborative manner on any textual dataset. As an example, you can try the Wikitext-103 dataset, composed of Wikipedia articles, widely used in natural language modeling, which you can download <a class='underline text-blue-400' target='_blank' href='https://dax-cdn.cdn.appdomain.cloud/dax-wikitext-103/1.0.1/wikitext-103.tar.gz'>here</a>. More information on how to connect the dataset at the next step."
                 },
-                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."',
-                sampleDatasetLink: 'https://dax-cdn.cdn.appdomain.cloud/dax-wikitext-103/1.0.1/wikitext-103.tar.gz'
+                model: 'The model follows the exact GPT-2 architecture and is implemented in TensorFlow.js. The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer. The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss. To accommodate all devices, the context length is currently kept at 128 and the batch size at 1.',
+                dataFormatInformation: 'You can use any natural language (text) dataset you like. For example the Wikitext-103 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 is: <i>"For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work ( except eight private performances for Ludwig II at Munich in 1884 and 1885 ) ."</i>',
+                sampleDatasetLink: 'https://dax-cdn.cdn.appdomain.cloud/dax-wikitext-103/1.0.1/wikitext-103.tar.gz',
+                sampleDatasetInstructions: 'Opening the link should start downloading a zip file. Unzip it and drag and drop the training set named "wiki.train.tokens" in the field below (or use the "Select File" button). Even though the file extension is ".tokens" it is indeed a text file. You can use "wiki.test.tokens" at the evaluation step after training a language model.'
             },
             trainingInformation: {
                 dataType: 'text',
-                modelID: 'wikitext-103-raw-model',
+                modelID: 'llm-raw-model',
                 preprocessingFunctions: [data.TextPreprocessing.Tokenize, data.TextPreprocessing.LeftPadding],
                 scheme: 'federated',
-                epochs: 5,
+                epochs: 6,
                 // Unused by wikitext because data already comes split
                 // But if set to 0 then the webapp doesn't display the validation metrics
                 validationSplit: 0.1,
                 roundDuration: 2,
-                batchSize: 1, // If set too high (e.g. 16) then firefox raises a WebGL error
+                batchSize: 1, // If set too high (e.g. 16) firefox raises a WebGL error
                 tokenizer: 'Xenova/gpt2',
                 maxSequenceLength: 128,
                 tensorBackend: 'gpt'

package/dist/index.d.ts

@@ -10,8 +10,9 @@ export { Logger, ConsoleLogger } from './logging/index.js';
 export { Memory, 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';
+export { Model, BatchLogs, EpochLogs, ValidationMetrics } from './models/index.js';
 export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
 export * from './types.js';
+export * as async_iterator from "./utils/async_iterator.js";

package/dist/index.js

@@ -10,8 +10,9 @@ export { ConsoleLogger } from './logging/index.js';
 export { Memory, 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';
+export { Model, EpochLogs } from './models/index.js';
 export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
 export * from './types.js';
+export * as async_iterator from "./utils/async_iterator.js";

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

@@ -3,5 +3,5 @@ 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 default function evaluate(model: tf.LayersModel, dataset: tf.data.Dataset<DataPoint>, maxEvalBatches: number): Promise<Record<'val_acc' | 'val_loss' | 'val_perplexity', number>>;
 export {};

package/dist/models/gpt/evaluate.js

@@ -38,7 +38,6 @@ export default async function evaluate(model, dataset, maxEvalBatches) {
     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

@@ -5,14 +5,15 @@ 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';
+import { BatchLogs, Model, EpochLogs } from "../index.js";
+import type { Prediction, Sample } from '../model.js';
+import { type GPTConfig } from './config.js';
 export type GPTSerialization = {
     weights: WeightsContainer;
     config?: GPTConfig;
 };
 export declare class GPT extends Model {
+    #private;
     private readonly model;
     constructor(partialConfig?: GPTConfig, layersModel?: tf.LayersModel);
     /**
@@ -24,7 +25,7 @@ export declare class GPT extends Model {
      * @param epochs the number of passes of the training dataset
      * @param tracker
      */
-    train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs, void>;
+    train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
     predict(input: Sample): Promise<Prediction>;
     generate(input: string, tokenizer: PreTrainedTokenizer, newTokens?: number): Promise<string>;
     get config(): Required<GPTConfig>;

package/dist/models/gpt/index.js

@@ -1,14 +1,20 @@
 /**
  * this code is taken from gpt-tfjs with modifications from @peacefulotter and @lukemovement
  **/
+import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../../index.js';
-import { Model } from '../model.js';
+import { Model, EpochLogs } from "../index.js";
 import { GPTForCausalLM } from './model.js';
+import { DEFAULT_CONFIG } from './config.js';
+import evaluate from './evaluate.js';
+import { List } from 'immutable';
 export class GPT extends Model {
     model;
+    #maxBatchCount;
     constructor(partialConfig, layersModel) {
         super();
         this.model = new GPTForCausalLM(partialConfig, layersModel);
+        this.#maxBatchCount = partialConfig?.maxIter ?? DEFAULT_CONFIG.maxIter;
     }
     /**
      * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs)
@@ -19,45 +25,65 @@ export class GPT extends Model {
      * @param epochs the number of passes of the training dataset
      * @param tracker
      */
-    async *train(trainingData, validationData, epochs = 1) {
+    async *train(trainingData, validationData) {
         this.model.compile();
+        const batches = await trainingData.iterator(); // tf.LazyIterator isn't an AsyncGenerator
+        let batchesLogs = List();
+        for (let batchNumber = 0; batchNumber < this.#maxBatchCount; batchNumber++) {
+            const iteration = await batches.next();
+            if (iteration.done)
+                break;
+            const batch = iteration.value;
+            const batchLogs = await this.#runBatch(batch);
+            tf.dispose(batch);
+            yield batchLogs;
+            batchesLogs = batchesLogs.push(batchLogs);
+        }
+        const validation = validationData && (await this.#evaluate(validationData));
+        return new EpochLogs(batchesLogs, validation);
+    }
+    async #runBatch(batch) {
         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; } },
+        await this.model.fitDataset(tf.data.array([batch]), {
+            epochs: 1,
+            verbose: 0, // don't pollute
+            callbacks: {
+                onEpochEnd: (_, cur) => {
+                    logs = cur;
+                },
+            },
+        });
+        if (logs === undefined)
+            throw new Error("batch didn't gave any logs");
+        const { loss, acc: accuracy } = logs;
+        if (loss === undefined || isNaN(loss))
+            throw new Error("training loss is undefined or NaN");
+        return {
+            accuracy,
+            loss,
+            memoryUsage: tf.memory().numBytes / 1024 / 1024 / 1024,
         };
-        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,
-                    accuracy: logs.acc
-                }
-            };
-            if (validationData !== undefined) {
-                if (val_loss === undefined || isNaN(val_loss) ||
-                    val_acc === undefined || isNaN(val_acc)) {
-                    throw new Error("Validation accuracy or loss is undefined or nan");
-                }
-                structuredLogs.validation = { accuracy: logs.val_acc, loss: logs.val_loss };
+    }
+    async #evaluate(dataset) {
+        const evaluation = await evaluate(this.model, dataset.map((t) => {
+            switch (t) {
+                case null:
+                case undefined:
+                    throw new Error("nullish value in dataset");
+                default:
+                    // TODO unsafe cast
+                    return t;
             }
-            yield structuredLogs;
-        }
+        }), this.config.maxEvalBatches);
+        return {
+            accuracy: evaluation.val_acc,
+            loss: evaluation.val_loss,
+        };
     }
     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');
+            throw new Error("prediction yield many Tensors but should have only returned one");
         }
         return Promise.resolve(ret);
     }

package/dist/models/gpt/model.js

@@ -37,29 +37,44 @@ class GPTModel extends tf.LayersModel {
         const evalDataset = trainingArgs.validationData;
         await callbacks.onTrainBegin?.();
         for (let epoch = 1; epoch <= trainingArgs.epochs; epoch++) {
+            let accuracyFraction = [0, 0];
             let averageLoss = 0;
-            let peakMemory = 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 = () => {
+                let preprocessingTime = performance.now();
+                await Promise.all([xs.data(), ys.data()]);
+                preprocessingTime = performance.now() - preprocessingTime;
+                // TODO include as a tensor inside the model
+                const accTensor = tf.tidy(() => {
                     const logits = this.apply(xs);
-                    if (Array.isArray(logits)) {
+                    if (Array.isArray(logits))
                         throw new Error('model outputs too many tensor');
-                    }
-                    if (logits instanceof tf.SymbolicTensor) {
+                    if (logits instanceof tf.SymbolicTensor)
                         throw new Error('model outputs symbolic tensor');
-                    }
-                    return tf.losses.softmaxCrossEntropy(ys, logits);
-                };
+                    return tf.metrics.categoricalAccuracy(ys, logits);
+                });
+                const accSize = accTensor.shape.reduce((l, r) => l * r, 1);
+                const accSumTensor = accTensor.sum();
+                const accSum = await accSumTensor.array();
+                tf.dispose(accSumTensor);
+                if (typeof accSum !== 'number')
+                    throw new Error('got multiple accuracy sum');
+                accuracyFraction = [accuracyFraction[0] + accSum, accuracyFraction[1] + accSize];
+                tf.dispose([accTensor]);
                 const lossTensor = tf.tidy(() => {
-                    const { grads, value: lossTensor } = this.optimizer.computeGradients(lossFn);
+                    const { grads, value: lossTensor } = this.optimizer.computeGradients(() => {
+                        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);
+                    });
                     const gradsClipped = clipByGlobalNormObj(grads, 1);
                     this.optimizer.applyGradients(gradsClipped);
                     return lossTensor;
@@ -75,9 +90,6 @@ class GPTModel extends tf.LayersModel {
                     console.log(iterationLogs);
                 }
                 const memory = tf.memory().numBytes / 1024 / 1024 / 1024;
-                if (memory > peakMemory) {
-                    peakMemory = memory;
-                }
                 console.log(`Epoch: ${epoch}`, `\tStep: ${iteration} / ${this.config.maxIter}`, `\tLoss: ${loss.toFixed(3)}`, `\tMemory: ${memory.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();
@@ -89,7 +101,7 @@ class GPTModel extends tf.LayersModel {
             }
             let logs = {
                 'loss': averageLoss / iteration,
-                'peakMemory': peakMemory
+                'acc': accuracyFraction[0] / accuracyFraction[1],
             };
             if (evalDataset !== undefined) {
                 logs = { ...logs, ...await evaluate(this, evalDataset, this.config.maxEvalBatches) };

package/dist/models/index.d.ts

@@ -1,4 +1,5 @@
-export { EpochLogs, Model } from './model.js';
+export { Model } from './model.js';
+export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js";
 export { GPT } from './gpt/index.js';
 export { GPTConfig } from './gpt/config.js';
 export { TFJS } from './tfjs.js';

package/dist/models/index.js

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

package/dist/models/logs.d.ts

@@ -0,0 +1,17 @@
+import { List } from "immutable";
+export interface ValidationMetrics {
+    accuracy: number;
+    loss: number;
+}
+export interface BatchLogs {
+    accuracy: number;
+    loss: number;
+    memoryUsage: number;
+}
+export declare class EpochLogs {
+    readonly validation?: ValidationMetrics | undefined;
+    readonly batches: List<BatchLogs>;
+    constructor(batches: Iterable<BatchLogs>, validation?: ValidationMetrics | undefined);
+    get training(): Record<"accuracy" | "loss", number>;
+    get peakMemory(): number;
+}

package/dist/models/logs.js

@@ -0,0 +1,22 @@
+import { List } from "immutable";
+export class EpochLogs {
+    validation;
+    batches;
+    constructor(batches, validation) {
+        this.validation = validation;
+        this.batches = List(batches);
+    }
+    get training() {
+        const sum = this.batches.reduce((acc, batch) => ({
+            accuracy: acc.accuracy + batch.accuracy,
+            loss: acc.loss + batch.loss,
+        }), { loss: 0, accuracy: 0 });
+        return {
+            accuracy: sum.accuracy / this.batches.size,
+            loss: sum.loss / this.batches.size,
+        };
+    }
+    get peakMemory() {
+        return this.batches.map((batch) => batch.memoryUsage).max() ?? 0;
+    }
+}

package/dist/models/model.d.ts

@@ -1,19 +1,7 @@
-/// <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;
-}
+import type { BatchLogs, EpochLogs } from "./logs.js";
 export type Prediction = tf.Tensor;
 export type Sample = tf.Tensor;
 /**
@@ -35,7 +23,7 @@ export declare abstract class Model implements Disposable {
      * @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>;
+    abstract train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
     /** Predict likely values */
     abstract predict(input: Sample): Promise<Prediction>;
     /**

package/dist/models/tfjs.d.ts

@@ -1,16 +1,18 @@
 import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../index.js';
-import { Model } from './index.js';
-import type { EpochLogs, Prediction, Sample } from './model.js';
 import type { Dataset } from '../dataset/index.js';
+import { BatchLogs, EpochLogs } from './index.js';
+import { Model } from './index.js';
+import type { Prediction, Sample } from './model.js';
 /** TensorFlow JavaScript model with standard training */
 export declare class TFJS extends Model {
+    #private;
     private readonly model;
     /** Wrap the given trainable model */
     constructor(model: tf.LayersModel);
     get weights(): WeightsContainer;
     set weights(ws: WeightsContainer);
-    train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs>;
+    train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
     predict(input: Sample): Promise<Prediction>;
     static deserialize(raw: tf.io.ModelArtifacts): Promise<Model>;
     serialize(): Promise<tf.io.ModelArtifacts>;

package/dist/models/tfjs.js

@@ -1,5 +1,7 @@
+import { List, Map } from 'immutable';
 import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../index.js';
+import { EpochLogs } from './index.js';
 import { Model } from './index.js';
 /** TensorFlow JavaScript model with standard training */
 export class TFJS extends Model {
@@ -18,50 +20,71 @@ export class TFJS extends Model {
     set weights(ws) {
         this.model.setWeights(ws.weights);
     }
-    async *train(trainingData, validationData, epochs = 1) {
-        for (let epoch = 0; epoch < epochs; epoch++) {
-            let logs;
-            let peakMemory = 0;
-            await this.model.fitDataset(trainingData, {
-                epochs: 1,
-                validationData,
-                callbacks: {
-                    onBatchEnd: (_) => {
-                        const currentMemory = tf.memory().numBytes / 1024 / 1024 / 1024; // GB
-                        if (currentMemory > peakMemory) {
-                            peakMemory = currentMemory;
-                        }
-                    },
-                    onEpochEnd: (_, cur) => { logs = cur; }
-                },
-            });
-            if (logs === undefined) {
-                throw new Error("Epoch didn't gave any logs");
-            }
-            const { loss, acc, val_acc, val_loss } = logs;
-            if (loss === undefined || isNaN(loss) || acc === undefined || isNaN(acc)) {
-                throw new Error("Training loss is undefined or nan");
-            }
-            const structuredLogs = {
-                epoch,
-                peakMemory,
-                training: {
-                    loss: logs.loss,
-                    accuracy: logs.acc,
-                }
+    async *train(trainingData, validationData) {
+        const batches = await trainingData.iterator(); // tf.LazyIterator isn't an AsyncGenerator
+        let batchesLogs = List();
+        for (let batchNumber = 0; true; batchNumber++) {
+            const iteration = await batches.next();
+            if (iteration.done)
+                break;
+            const batch = iteration.value;
+            const batchLogs = {
+                batch: batchNumber,
+                ...(await this.#runBatch(batch)),
             };
-            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;
+            tf.dispose(batch);
+            yield batchLogs;
+            batchesLogs = batchesLogs.push(batchLogs);
         }
+        const validation = validationData && (await this.#evaluate(validationData));
+        return new EpochLogs(batchesLogs, validation);
+    }
+    async #runBatch(batch) {
+        let logs;
+        await this.model.fitDataset(tf.data.array([batch]), {
+            epochs: 1,
+            verbose: 0, // don't pollute
+            callbacks: {
+                onEpochEnd: (_, cur) => {
+                    logs = cur;
+                },
+            },
+        });
+        if (logs === undefined)
+            throw new Error("batch didn't gave any logs");
+        const { loss, acc: accuracy } = logs;
+        if (loss === undefined || isNaN(loss))
+            throw new Error("training loss is undefined or NaN");
+        return {
+            accuracy,
+            loss,
+            memoryUsage: tf.memory().numBytes / 1024 / 1024 / 1024,
+        };
+    }
+    async #evaluate(dataset) {
+        const evaluation = await this.model.evaluateDataset(dataset.map((t) => {
+            switch (t) {
+                case null:
+                case undefined:
+                    throw new Error("nullish value in dataset");
+                default:
+                    return t;
+            }
+        }));
+        const metricToValue = Map(List(this.model.metricsNames).zip(Array.isArray(evaluation)
+            ? List(await Promise.all(evaluation.map((t) => t.data())))
+            : List.of(await evaluation.data()))).map((values) => {
+            if (values.length !== 1)
+                throw new Error("more than one metric value");
+            return values[0];
+        });
+        const [accuracy, loss] = [
+            metricToValue.get("acc"),
+            metricToValue.get("loss"),
+        ];
+        if (accuracy === undefined || loss === undefined)
+            throw new Error("some needed metrics are missing");
+        return { accuracy, loss };
     }
     predict(input) {
         const ret = this.model.predict(input);

package/dist/task/display_information.d.ts

@@ -10,5 +10,6 @@ export interface DisplayInformation {
     headers?: string[];
     dataExampleImage?: string;
     sampleDatasetLink?: string;
+    sampleDatasetInstructions?: string;
 }
 export declare function isDisplayInformation(raw: unknown): raw is DisplayInformation;

package/dist/task/display_information.js

@@ -4,13 +4,14 @@ export function isDisplayInformation(raw) {
     if (typeof raw !== 'object' || raw === null) {
         return false;
     }
-    const { dataExample, dataExampleImage, dataExampleText, dataFormatInformation, sampleDatasetLink, headers, model, summary, taskTitle, } = raw;
+    const { dataExample, dataExampleImage, dataExampleText, dataFormatInformation, sampleDatasetLink, sampleDatasetInstructions, headers, model, summary, taskTitle, } = raw;
     if (typeof taskTitle !== 'string' ||
         (dataExampleText !== undefined && typeof dataExampleText !== 'string') ||
         (sampleDatasetLink !== undefined && typeof sampleDatasetLink !== 'string') ||
         (dataFormatInformation !== undefined && typeof dataFormatInformation !== 'string') ||
         (model !== undefined && typeof model !== 'string') ||
-        (dataExampleImage !== undefined && typeof dataExampleImage !== 'string')) {
+        (dataExampleImage !== undefined && typeof dataExampleImage !== 'string') ||
+        (sampleDatasetInstructions !== undefined && typeof sampleDatasetInstructions !== 'string')) {
         return false;
     }
     if (!isSummary(summary)) {
@@ -46,6 +47,7 @@ export function isDisplayInformation(raw) {
         dataExampleText,
         dataFormatInformation,
         sampleDatasetLink,
+        sampleDatasetInstructions,
         headers,
         model,
         summary,

package/dist/training/disco.d.ts

@@ -1,4 +1,4 @@
-import type { data, Logger, Memory, Task, TrainingInformation } from '../index.js';
+import { BatchLogs, data, EpochLogs, Logger, Memory, Task, TrainingInformation } from '../index.js';
 import { client as clients } from '../index.js';
 import type { Aggregator } from '../aggregator/index.js';
 import type { RoundLogs } from './trainer/trainer.js';
@@ -22,13 +22,25 @@ export declare class Disco {
     private readonly client;
     private readonly trainer;
     constructor(task: Task, options: DiscoOptions);
-    /**
-     * Starts a training instance for the Disco object's task on the provided data tuple.
-     * @param dataTuple The data tuple
-     */
-    fit(dataTuple: data.DataSplit): AsyncGenerator<RoundLogs & {
+    /** Train on dataset, yielding logs of every round. */
+    trainByRound(dataTuple: data.DataSplit): AsyncGenerator<RoundLogs & {
         participants: number;
     }>;
+    /** Train on dataset, yielding logs of every epoch. */
+    trainByEpoch(dataTuple: data.DataSplit): AsyncGenerator<EpochLogs>;
+    /** Train on dataset, yielding logs of every batch. */
+    trainByBatch(dataTuple: data.DataSplit): AsyncGenerator<BatchLogs>;
+    /** Run whole train on dataset. */
+    trainFully(dataTuple: data.DataSplit): Promise<void>;
+    /**
+    * Train on dataset, yield the nested steps.
+    *
+    * Don't forget to await the yielded generator otherwise nothing will progress.
+    * If you don't care about the whole process, use one of the other train methods.
+    **/
+    train(dataTuple: data.DataSplit): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs & {
+        participants: number;
+    }>>;
     /**
      * Stops the ongoing training instance without disconnecting the client.
      */

package/dist/training/disco.js

@@ -1,5 +1,8 @@
+import { List } from 'immutable';
+import { async_iterator } from '../index.js';
 import { client as clients, EmptyMemory, ConsoleLogger } from '../index.js';
 import { MeanAggregator } from '../aggregator/mean.js';
+import { enumerate, split } from '../utils/async_iterator.js';
 import { TrainerBuilder } from './trainer/trainer_builder.js';
 /**
  * Top-level class handling distributed training from a client's perspective. It is meant to be
@@ -58,35 +61,79 @@ export class Disco {
         const trainerBuilder = new TrainerBuilder(this.memory, this.task);
         this.trainer = trainerBuilder.build(this.client, options.scheme !== 'local');
     }
+    /** Train on dataset, yielding logs of every round. */
+    async *trainByRound(dataTuple) {
+        for await (const round of this.train(dataTuple)) {
+            const [roundGen, roundLogs] = async_iterator.split(round);
+            for await (const epoch of roundGen)
+                for await (const _ of epoch)
+                    ;
+            yield await roundLogs;
+        }
+    }
+    /** Train on dataset, yielding logs of every epoch. */
+    async *trainByEpoch(dataTuple) {
+        for await (const round of this.train(dataTuple)) {
+            for await (const epoch of round) {
+                const [epochGen, epochLogs] = async_iterator.split(epoch);
+                for await (const _ of epochGen)
+                    ;
+                yield await epochLogs;
+            }
+        }
+    }
+    /** Train on dataset, yielding logs of every batch. */
+    async *trainByBatch(dataTuple) {
+        for await (const round of this.train(dataTuple))
+            for await (const epoch of round)
+                yield* epoch;
+    }
+    /** Run whole train on dataset. */
+    async trainFully(dataTuple) {
+        for await (const round of this.train(dataTuple))
+            for await (const epoch of round)
+                for await (const _ of epoch)
+                    ;
+    }
     /**
-     * Starts a training instance for the Disco object's task on the provided data tuple.
-     * @param dataTuple The data tuple
-     */
+    * Train on dataset, yield the nested steps.
+    *
+    * Don't forget to await the yielded generator otherwise nothing will progress.
+    * If you don't care about the whole process, use one of the other train methods.
+    **/
     // TODO RoundLogs should contain number of participants but Trainer doesn't need client
-    async *fit(dataTuple) {
+    async *train(dataTuple) {
         this.logger.success("Training started.");
         const trainData = dataTuple.train.preprocess().batch();
         const validationData = dataTuple.validation?.preprocess().batch() ?? trainData;
         await this.client.connect();
         const trainer = await this.trainer;
-        for await (const roundLogs of trainer.fitModel(trainData.dataset, validationData.dataset)) {
-            let msg = `Round: ${roundLogs.round}\n`;
-            for (const epochLogs of roundLogs.epochs.values()) {
-                msg += `    Epoch: ${epochLogs.epoch}\n`;
-                msg += `        Training loss: ${epochLogs.training.loss}\n`;
-                if (epochLogs.training.accuracy !== undefined) {
-                    msg += `        Training accuracy: ${epochLogs.training.accuracy}\n`;
+        for await (const [round, epochs] of enumerate(trainer.fitModel(trainData.dataset, validationData.dataset))) {
+            yield async function* () {
+                let epochsLogs = List();
+                for await (const [epoch, batches] of enumerate(epochs)) {
+                    const [gen, returnedEpochLogs] = split(batches);
+                    yield gen;
+                    const epochLogs = await returnedEpochLogs;
+                    epochsLogs = epochsLogs.push(epochLogs);
+                    this.logger.success([
+                        `Round: ${round}`,
+                        `  Epoch: ${epoch}`,
+                        `    Training loss: ${epochLogs.training.loss}`,
+                        `    Training accuracy: ${epochLogs.training.accuracy}`,
+                        epochLogs.validation !== undefined
+                            ? `    Validation loss: ${epochLogs.validation.loss}`
+                            : "",
+                        epochLogs.validation !== undefined
+                            ? `    Validation accuracy: ${epochLogs.validation.accuracy}`
+                            : "",
+                    ].join("\n"));
                 }
-                if (epochLogs.validation !== undefined) {
-                    msg += `        Validation loss: ${epochLogs.validation.loss}\n`;
-                    msg += `        Validation accuracy: ${epochLogs.validation.accuracy}\n`;
-                }
-            }
-            this.logger.success(msg);
-            yield {
-                ...roundLogs,
-                participants: this.client.nodes.size + 1 // add ourself
-            };
+                return {
+                    epochs: epochsLogs,
+                    participants: this.client.nodes.size + 1, // add ourself
+                };
+            }.bind(this)();
         }
         this.logger.success("Training finished.");
     }

package/dist/training/trainer/trainer.d.ts

@@ -1,9 +1,8 @@
 import type tf from "@tensorflow/tfjs";
 import { List } from "immutable";
 import type { Model, Task } from "../../index.js";
-import { EpochLogs } from "../../models/model.js";
+import { BatchLogs, EpochLogs } from "../../models/index.js";
 export interface RoundLogs {
-    round: number;
     epochs: List<EpochLogs>;
 }
 /** Abstract class whose role is to train a model with a given dataset. This can be either done
@@ -29,5 +28,5 @@ export declare abstract class Trainer {
      * Start training the model with the given dataset
      * @param dataset
      */
-    fitModel(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<RoundLogs>;
+    fitModel(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs>, void>;
 }

package/dist/training/trainer/trainer.js

@@ -1,4 +1,5 @@
 import { List } from "immutable";
+import * as async_iterator from "../../utils/async_iterator.js";
 /** Abstract class whose role is to train a model with a given dataset. This can be either done
  * locally (alone) or in a distributed way with collaborators.
  *
@@ -16,6 +17,8 @@ export class Trainer {
         this.model = model;
         this.#roundDuration = task.trainingInformation.roundDuration;
         this.#epochs = task.trainingInformation.epochs;
+        if (!Number.isInteger(this.#epochs / this.#roundDuration))
+            throw new Error(`round duration doesn't divide epochs`);
     }
     /**
      * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
@@ -28,25 +31,31 @@ export class Trainer {
      * @param dataset
      */
     async *fitModel(dataset, valDataset) {
-        if (this.training !== undefined) {
+        if (this.training !== undefined)
             throw new Error("training already running, cancel it before launching a new one");
+        try {
+            this.training = this.#runRounds(dataset, valDataset);
+            yield* this.training;
         }
-        await this.onRoundBegin(0);
-        this.training = this.model.train(dataset, valDataset, this.#epochs);
-        for await (const logs of this.training) {
-            // for now, round (sharing on network) == epoch (full pass over local data)
-            yield {
-                round: logs.epoch,
-                epochs: List.of(logs),
-            };
-            if (logs.epoch % this.#roundDuration === 0) {
-                const round = Math.trunc(logs.epoch / this.#roundDuration);
-                await this.onRoundEnd(round);
-                await this.onRoundBegin(round);
-            }
+        finally {
+            this.training = undefined;
         }
-        const round = Math.trunc(this.#epochs / this.#roundDuration);
-        await this.onRoundEnd(round);
-        this.training = undefined;
+    }
+    async *#runRounds(dataset, valDataset) {
+        const totalRound = Math.trunc(this.#epochs / this.#roundDuration);
+        for (let round = 0; round < totalRound; round++) {
+            await this.onRoundBegin(round);
+            yield this.#runRound(dataset, valDataset);
+            await this.onRoundEnd(round);
+        }
+    }
+    async *#runRound(dataset, valDataset) {
+        let epochsLogs = List();
+        for (let epoch = 0; epoch < this.#roundDuration; epoch++) {
+            const [gen, epochLogs] = async_iterator.split(this.model.train(dataset, valDataset));
+            yield gen;
+            epochsLogs = epochsLogs.push(await epochLogs);
+        }
+        return { epochs: epochsLogs };
     }
 }

package/dist/utils/async_iterator.d.ts

@@ -0,0 +1,11 @@
+import { List } from "immutable";
+/**
+ * Split yields from return value
+ *
+ * You need to consume the iterator to resolve the returned value
+ **/
+export declare function split<T, U>(iter: AsyncIterator<T, U>): [AsyncGenerator<T, U>, Promise<U>];
+/** Zip iterator with a infinite counter */
+export declare function enumerate<T, U>(iter: AsyncIterator<T, U> | Iterator<T, U>): AsyncGenerator<[number, T], U>;
+/** Run the whole iterator to get yielded & returned */
+export declare function gather<T, U>(iter: AsyncIterator<T, U>): Promise<[List<T>, U]>;

package/dist/utils/async_iterator.js

@@ -0,0 +1,63 @@
+import { List } from "immutable";
+// `Promise.withResolvers` not widely deployed
+function PromiseWithResolvers() {
+    let resolve, reject;
+    resolve = reject = () => {
+        // should not happen as Promise are run on creation
+        throw new Error("race condition triggered");
+    };
+    const promise = new Promise((res, rej) => {
+        resolve = res;
+        reject = rej;
+    });
+    return [promise, resolve, reject];
+}
+/**
+ * Split yields from return value
+ *
+ * You need to consume the iterator to resolve the returned value
+ **/
+export function split(iter) {
+    const [returnPromise, returnResolve, returnReject] = PromiseWithResolvers();
+    return [
+        (async function* () {
+            try {
+                while (true) {
+                    const v = await iter.next();
+                    if (!v.done) {
+                        yield v.value;
+                        continue;
+                    }
+                    returnResolve(v.value);
+                    return v.value;
+                }
+            }
+            catch (e) {
+                returnReject(e);
+                throw e;
+            }
+        })(),
+        returnPromise,
+    ];
+}
+/** Zip iterator with a infinite counter */
+export function enumerate(iter) {
+    return (async function* () {
+        for (let i = 0;; i++) {
+            const v = await iter.next();
+            if (v.done)
+                return v.value;
+            yield [i, v.value];
+        }
+    })();
+}
+/** Run the whole iterator to get yielded & returned */
+export async function gather(iter) {
+    let elems = List();
+    for (;;) {
+        const v = await iter.next();
+        if (v.done)
+            return [elems, v.value];
+        elems = elems.push(v.value);
+    }
+}

package/package.json

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

package/dist/default_tasks/skin_condition.d.ts

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

package/dist/default_tasks/skin_condition.js

@@ -1,80 +0,0 @@
-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,
-                tensorBackend: 'tfjs'
-            }
-        };
-    },
-    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));
-    }
-};