*: framework agnostic preprocessing

cli/src/args.ts

@@ -1,11 +1,11 @@
 import { parse } from 'ts-command-line-args'
 import { Map, Set } from 'immutable'
 
-import type { TaskProvider } from '@epfml/discojs'
+import type { DataType, TaskProvider } from "@epfml/discojs";
 import { defaultTasks } from '@epfml/discojs'
 
 interface BenchmarkArguments {
-  provider: TaskProvider
+  provider: TaskProvider<DataType>
   numberOfUsers: number
   epochs: number
   roundDuration: number
@@ -37,7 +37,7 @@ const unsafeArgs = parse<BenchmarkUnsafeArguments>(
 )
 
 const supportedTasks = Map(
-  Set.of(
+  Set.of<TaskProvider<"image"> | TaskProvider<"tabular">>(
     defaultTasks.cifar10,
     defaultTasks.lusCovid,
     defaultTasks.simpleFace,
@@ -67,6 +67,6 @@ export const args: BenchmarkArguments = {
 
       return task;
     },
-    getModel: provider.getModel,
+    getModel: () => provider.getModel(),
   },
 };

cli/src/benchmark_gpt.ts

@@ -1,8 +1,8 @@
+import { List } from "immutable";
 import { parse } from "ts-command-line-args";
-import * as tf from "@tensorflow/tfjs"
 import { AutoTokenizer } from "@xenova/transformers";
 
-import { fetchTasks, models, async_iterator, defaultTasks, processing } from "@epfml/discojs";
+import { fetchTasks, models, async_iterator, defaultTasks, processing, Task } from "@epfml/discojs";
 import { loadModelFromDisk, loadText } from '@epfml/discojs-node'
 
 import { Server } from "server";
@@ -41,15 +41,6 @@ const args = { ...defaultArgs, ...parsedArgs }
  * Benchmark results are reported in https://github.com/epfml/disco/pull/659
  */
 
-function intoTFGenerator<T extends tf.TensorContainer>(
-  iter: AsyncIterable<T>,
-): tf.data.Dataset<T> {
-  // @ts-expect-error generator
-  return tf.data.generator(async function* () {
-    yield* iter;
-  });
-}
-
 async function main(args: Required<CLIArguments>): Promise<void> { 
   const { inference: benchmarkInference, modelType,
     contextLength, batchSize, modelPath } = args
@@ -59,7 +50,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
 
   // Fetch the wikitext task from the server
   const tasks = await fetchTasks(url)
-  const task = tasks.get('llm_task')
+  const task = tasks.get('llm_task') as Task<'text'> | undefined
   if (task === undefined) { throw new Error('task not found') }  
 
   const tokenizerName = task.trainingInformation.tokenizer
@@ -89,32 +80,10 @@ async function main(args: Required<CLIArguments>): Promise<void> {
     const dataset = loadText('../datasets/wikitext/wiki.train.tokens')
 
     const maxLength = task.trainingInformation.maxSequenceLength ?? (tokenizer.model_max_length as number) + 1
-    // TODO will be easier when preproccessing is redone
-    const preprocessedDataset = intoTFGenerator(
-      dataset
-        .map((line) =>
-          processing.tokenizeAndLeftPad(line, tokenizer, maxLength),
-        )
-        .batch(batchSize)
-        .map((batch) =>
-          tf.tidy(() => ({
-            xs: tf.tensor2d(
-              batch.map((tokens) => tokens.slice(0, -1)).toArray(),
-            ),
-            ys: tf.stack(
-              batch
-                .map(
-                  (tokens) =>
-                    tf.oneHot(
-                      tokens.slice(1),
-                      tokenizer.model.vocab.length + 1,
-                    ) as tf.Tensor2D,
-                )
-                .toArray(),
-            ) as tf.Tensor3D,
-          })),
-        ),
-    );
+    const preprocessedDataset = dataset
+      .map((line) => processing.tokenizeAndLeftPad(line, tokenizer, maxLength))
+      .map((tokens) => [tokens.pop(), tokens.last()] as [List<number>, number])
+      .batch(batchSize);
 
     // Init and train the model
     const model = new models.GPT(config)
@@ -143,13 +112,20 @@ async function main(args: Required<CLIArguments>): Promise<void> {
     const iterations = 10
     console.log("Generating", nbNewTokens, "new tokens")
 
+    let tokens = List(
+      (tokenizer(prompt, { return_tensor: false }) as { input_ids: number[] })
+        .input_ids,
+    );
+
     let inferenceTime = 0
     for (let i = 0; i < iterations; i++) {
       const timeStart = performance.now()
-      const _ = await model.generate(prompt, tokenizer, nbNewTokens)
+      for (let n = 0; n < nbNewTokens; n++) {
+        const next: number = (await model.predict(List.of(tokens))).first();
+	tokens = tokens.push(next)
+      }
       inferenceTime += performance.now() - timeStart
     }
-    // Overall average includes tokenization, token sampling and de-tokenization
     console.log(`Inference time: ${(inferenceTime/ nbNewTokens / iterations).toFixed(2)} ms/token`)
   }
   await new Promise((resolve, reject) => {

cli/src/cli.ts

@@ -4,7 +4,14 @@ import "@tensorflow/tfjs-node"
 import { List, Range } from 'immutable'
 import fs from 'node:fs/promises'
 
-import type { RoundLogs, Task, TaskProvider, TypedLabeledDataset } from '@epfml/discojs'
+import type {
+  Dataset,
+  DataFormat,
+  DataType,
+  RoundLogs,
+  Task,
+  TaskProvider,
+} from "@epfml/discojs";
 import { Disco, aggregator as aggregators, client as clients } from '@epfml/discojs'
 import { Server } from 'server'
 
@@ -18,10 +25,10 @@ async function arrayFromAsync<T>(iter: AsyncIterable<T>): Promise<T[]> {
   return ret;
 }
 
-async function runUser(
-  task: Task,
+async function runUser<D extends DataType>(
+  task: Task<D>,
   url: URL,
-  data: TypedLabeledDataset,
+  data: Dataset<DataFormat.Raw[D]>,
 ): Promise<List<RoundLogs>> {
   const trainingScheme = task.trainingInformation.scheme
   const aggregator = aggregators.getAggregator(task)
@@ -34,7 +41,10 @@ async function runUser(
   return logs;
 }
 
-async function main (provider: TaskProvider, numberOfUsers: number): Promise<void> {
+async function main<D extends DataType>(
+  provider: TaskProvider<D>,
+  numberOfUsers: number,
+): Promise<void> {
   const task = provider.getTask()
   console.log(`Started ${task.trainingInformation.scheme} training of ${task.id}`)
   console.log({ args })

cli/src/data.ts

@@ -1,10 +1,16 @@
 import path from "node:path";
 
-import type { Dataset, Image, Task, TypedLabeledDataset } from "@epfml/discojs";
+import type {
+  Dataset,
+  DataFormat,
+  DataType,
+  Image,
+  Task,
+} from "@epfml/discojs";
 import { loadCSV, loadImagesInDir } from "@epfml/discojs-node";
 import { Repeat } from "immutable";
 
-async function loadSimpleFaceData(): Promise<Dataset<[Image, string]>> {
+async function loadSimpleFaceData(): Promise<Dataset<DataFormat.Raw["image"]>> {
   const folder = path.join("..", "datasets", "simple_face");
 
   const [adults, childs]: Dataset<[Image, string]>[] = [
@@ -15,7 +21,7 @@ async function loadSimpleFaceData(): Promise<Dataset<[Image, string]>> {
   return adults.chain(childs);
 }
 
-async function loadLusCovidData(): Promise<Dataset<[Image, string]>> {
+async function loadLusCovidData(): Promise<Dataset<DataFormat.Raw["image"]>> {
   const folder = path.join("..", "datasets", "lus_covid");
 
   const [positive, negative]: Dataset<[Image, string]>[] = [
@@ -30,24 +36,22 @@ async function loadLusCovidData(): Promise<Dataset<[Image, string]>> {
   return positive.chain(negative);
 }
 
-export async function getTaskData(task: Task): Promise<TypedLabeledDataset> {
+export async function getTaskData<D extends DataType>(
+  task: Task<D>,
+): Promise<Dataset<DataFormat.Raw[D]>> {
   switch (task.id) {
     case "simple_face":
-      return ["image", await loadSimpleFaceData()];
+      return (await loadSimpleFaceData()) as Dataset<DataFormat.Raw[D]>;
     case "titanic":
-      return [
-        "tabular",
-        loadCSV(path.join("..", "datasets", "titanic_train.csv")),
-      ];
+      return loadCSV(
+        path.join("..", "datasets", "titanic_train.csv"),
+      ) as Dataset<DataFormat.Raw[D]>;
     case "cifar10":
-      return [
-        "image",
-        (await loadImagesInDir(path.join("..", "datasets", "CIFAR10"))).zip(
-          Repeat("cat"),
-        ),
-      ];
+      return (
+        await loadImagesInDir(path.join("..", "datasets", "CIFAR10"))
+      ).zip(Repeat("cat")) as Dataset<DataFormat.Raw[D]>;
     case "lus_covid":
-      return ["image", await loadLusCovidData()];
+      return (await loadLusCovidData()) as Dataset<DataFormat.Raw[D]>;
     default:
       throw new Error(`Data loader for ${task.id} not implemented.`);
   }

discojs-node/src/loaders/csv.ts

@@ -24,7 +24,7 @@ export function load(path: string): Dataset<Partial<Record<string, string>>> {
 
     for await (const row of stream) {
       if (!isRecordOfString(row))
-        throw new Error("excepted object of string to string");
+        throw new Error("expected object of string to string");
       yield row;
     }
   });

discojs-node/src/model_loader.ts

@@ -1,17 +1,23 @@
-import fs from 'node:fs/promises'
-import { serialization, models } from '@epfml/discojs'
+import fs from "node:fs/promises";
 
-export async function saveModelToDisk(model: models.Model, modelFolder: string, modelFileName: string): Promise<void> {
-  try {
-    await fs.access(modelFolder)
-  } catch {
-    await fs.mkdir(modelFolder)
-  }
-  const encoded = await serialization.model.encode(model)
-  await fs.writeFile(`${modelFolder}/${modelFileName}`, encoded)
+import type { models, DataType } from "@epfml/discojs";
+import { serialization } from "@epfml/discojs";
+
+export async function saveModelToDisk(
+  model: models.Model<DataType>,
+  modelFolder: string,
+  modelFileName: string,
+): Promise<void> {
+  const encoded = await serialization.model.encode(model);
+
+  await fs.mkdir(modelFolder, { recursive: true });
+  await fs.writeFile(`${modelFolder}/${modelFileName}`, encoded);
 }
 
-export async function loadModelFromDisk(modelPath: string): Promise<models.Model> {
-  const content = await fs.readFile(modelPath)
-  return await serialization.model.decode(content) as models.GPT
+export async function loadModelFromDisk(
+  modelPath: string,
+): Promise<models.Model<DataType>> {
+  const content = await fs.readFile(modelPath);
+
+  return await serialization.model.decode(content);
 }

discojs-web/package.json

@@ -27,6 +27,6 @@
     "@types/papaparse": "5",
     "jsdom": "25",
     "nodemon": "3",
-    "vitest": "1"
+    "vitest": "2"
   }
 }

discojs-web/src/loaders/csv.ts

@@ -30,7 +30,7 @@ export function load(file: File): Dataset<Partial<Record<string, string>>> {
 
           const rows = results.data.map((row) => {
             if (!isRecordOfString(row))
-              throw new Error("excepted object of string to string");
+              throw new Error("expected object of string to string");
 
             return row;
           });

discojs/package.json

@@ -19,6 +19,8 @@
   },
   "homepage": "https://github.com/epfml/disco#readme",
   "dependencies": {
+    "@jimp/core": "1",
+    "@jimp/plugin-resize": "1",
     "@msgpack/msgpack": "^3.0.0-beta2",
     "@tensorflow/tfjs": "4",
     "@xenova/transformers": "2",
@@ -31,7 +33,7 @@
   },
   "devDependencies": {
     "@tensorflow/tfjs-node": "4",
-    "@types/chai": "4",
+    "@types/chai": "5",
     "@types/mocha": "10",
     "@types/simple-peer": "9",
     "chai": "5",

discojs/src/aggregator/get.ts

@@ -1,8 +1,8 @@
-import type { Task } from '../index.js'
+import type { DataType, Task } from '../index.js'
 import { aggregator } from '../index.js'
 
 type AggregatorOptions = Partial<{
-  scheme: Task['trainingInformation']['scheme'], // if undefined, fallback on task.trainingInformation.scheme
+  scheme: Task<DataType>["trainingInformation"]["scheme"]; // if undefined, fallback on task.trainingInformation.scheme
   roundCutOff: number, // MeanAggregator
   threshold: number, // MeanAggregator
   thresholdType: 'relative' | 'absolute', // MeanAggregator
@@ -26,7 +26,10 @@ type AggregatorOptions = Partial<{
  * @param options Options passed down to the aggregator's constructor
  * @returns The aggregator
  */
-export function getAggregator(task: Task, options: AggregatorOptions = {}): aggregator.Aggregator {
+export function getAggregator(
+  task: Task<DataType>,
+  options: AggregatorOptions = {},
+): aggregator.Aggregator {
   const aggregationStrategy = task.trainingInformation.aggregationStrategy ?? 'mean'
   const scheme = options.scheme ?? task.trainingInformation.scheme
 

discojs/src/client/client.ts

@@ -1,6 +1,12 @@
 import createDebug from "debug";
 
-import type { Model, Task, WeightsContainer, RoundStatus } from '../index.js'
+import type {
+  DataType,
+  Model,
+  RoundStatus,
+  Task,
+  WeightsContainer,
+} from "../index.js";
 import { serialization } from '../index.js'
 import type { NodeID } from './types.js'
 import type { EventConnection } from './event_connection.js'
@@ -38,7 +44,7 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
 
   constructor (
     public readonly url: URL, // The network server's URL to connect to
-    public readonly task: Task, // The client's corresponding task
+    public readonly task: Task<DataType>, // The client's corresponding task
     public readonly aggregator: Aggregator,
   ) {
     super()
@@ -61,7 +67,7 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
    * This method is overriden by the federated and decentralized clients
    * By default, it fetches and returns the server's base model
    */
-  async connect(): Promise<Model> {
+  async connect(): Promise<Model<DataType>> {
     return this.getLatestModel()
   }
 
@@ -164,7 +170,7 @@ export abstract class Client extends EventEmitter<{'status': RoundStatus}>{
    * Fetches the latest model available on the network's server, for the adequate task.
    * @returns The latest model
    */
-  async getLatestModel (): Promise<Model> {
+  async getLatestModel (): Promise<Model<DataType>> {
     const url = new URL('', this.url.href)
     if (!url.pathname.endsWith('/')) {
       url.pathname += '/'

discojs/src/client/decentralized/decentralized_client.ts

@@ -1,7 +1,7 @@
 import createDebug from "debug";
 import { Map, Set } from 'immutable'
 
-import type { Model, WeightsContainer } from "../../index.js";
+import type { DataType, Model, WeightsContainer } from "../../index.js";
 import { serialization } from "../../index.js";
 import { Client,  shortenId } from '../client.js'
 import { type NodeID } from '../index.js'
@@ -44,7 +44,7 @@ export class DecentralizedClient extends Client {
    * create peer-to-peer WebRTC connections with peers. The server is used to exchange
    * peers network information.
    */
-  override async connect(): Promise<Model> {
+  override async connect(): Promise<Model<DataType>> {
     const model = await super.connect()  // Get the server base model
     const serverURL = new URL('', this.url.href)
     switch (this.url.protocol) {