Expand api surface

.gitignore

@@ -6,6 +6,7 @@
 # stored trained models
 /server/models/
 /models/
+/docs/examples/models/
 
 # tsc built
 /*/dist/

.knip.json

@@ -5,12 +5,12 @@
       "ignoreDependencies": ["typescript-eslint", "supports-color"]
     },
     "webapp": {
-      "ignore": ["env.d.ts"],
       "ignoreDependencies": ["buffer"]
     },
     "isomorphic-wrtc": {},
     "cli": {
       "entry": [
+        "src/cli.ts",
         "src/benchmark_gpt.ts",
         "src/hellaswag_gpt.ts",
         "src/train_gpt.ts"

Dockerfile

@@ -18,6 +18,7 @@ RUN pnpm ci
 COPY discojs/ discojs/
 COPY discojs-node/ discojs-node/
 COPY tsconfig.base.json .
+COPY tsconfig.base.lib.json .
 RUN pnpm -F discojs -F discojs-node run build
 
 COPY server/ server/

cli/src/args.ts

@@ -279,6 +279,6 @@ export const args: BenchmarkArguments = {
 
       return task;
     },
-    getModel: () => provider.getModel(),
+    modelCard: provider.modelCard,
   },
 };

cli/src/benchmark_gpt.ts

@@ -6,6 +6,7 @@ import type { Network, Task } from "@epfml/discojs";
 import {
   async_iterator,
   defaultTasks,
+  defaultModels,
   fetchTasks,
   models,
 } from "@epfml/discojs";
@@ -86,7 +87,10 @@ async function main(args: Required<CLIArguments>): Promise<void> {
   } = args;
 
   // Launch a server instance
-  const server = await Server.with(defaultTasks.wikitext);
+  const server = await Server.with(
+    [defaultModels.Wikitext],
+    [defaultTasks.wikitext],
+  );
   const [handle, url] = await server.serve();
 
   // Fetch the wikitext task from the server

cli/tsconfig.json

@@ -1,14 +1,14 @@
 {
-  "extends": "../tsconfig.base.json",
+  "extends": "../tsconfig.base.lib.json",
   "references": [
     {
-      "path": "../discojs"
+      "path": "../discojs/tsconfig.lib.json"
     },
     {
-      "path": "../discojs-node"
+      "path": "../discojs-node/tsconfig.lib.json"
     },
     {
-      "path": "../server"
+      "path": "../server/tsconfig.lib.json"
     }
   ],
   "compilerOptions": {

discojs-node/package.json

@@ -2,11 +2,15 @@
   "name": "@epfml/discojs-node",
   "version": "3.0.0",
   "type": "module",
-  "exports": "./dist/index.js",
+  "exports": {
+    "@disco/source": "./src/index.ts",
+    "default": "./dist/index.js",
+    "types": "./dist/index.d.ts"
+  },
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs/dist --watch . --exec pnpm run",
-    "build": "tsc --build",
+    "build": "tsc --build tsconfig.lib.json",
     "test": "cd .. && vitest --run --project=discojs-node"
   },
   "repository": {

discojs-node/tsconfig.lib.json

@@ -1,11 +1,15 @@
 {
-  "extends": "../tsconfig.base.json",
+  "extends": "../tsconfig.base.lib.json",
   "references": [
     {
-      "path": "../discojs"
+      "path": "../discojs/tsconfig.lib.json"
     }
   ],
-  "compilerOptions": { "rootDir": "./src", "outDir": "dist" },
+  "compilerOptions": {
+    "rootDir": "./src",
+    "outDir": "dist",
+    "composite": true
+  },
   "include": ["src"],
   "exclude": ["**/*.spec.ts"]
 }

discojs-node/tsconfig.vitest.json

@@ -2,7 +2,7 @@
   "extends": "../tsconfig.base.json",
   "references": [
     {
-      "path": "../discojs"
+      "path": "../discojs/tsconfig.vitest.json"
     }
   ],
   "compilerOptions": { "noEmit": true },

discojs-web/package.json

@@ -2,11 +2,15 @@
   "name": "@epfml/discojs-web",
   "version": "3.0.0",
   "type": "module",
-  "exports": "./dist/index.js",
+  "exports": {
+    "@disco/source": "./src/index.ts",
+    "default": "./dist/index.js",
+    "types": "./dist/index.d.ts"
+  },
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs/dist --watch . --exec pnpm run",
-    "build": "tsc --build",
+    "build": "tsc --build tsconfig.lib.json",
     "test": "cd .. && vitest --run --project=discojs-web"
   },
   "repository": {

discojs-web/tsconfig.lib.json

@@ -1,14 +1,15 @@
 {
-  "extends": "../tsconfig.base.json",
+  "extends": "../tsconfig.base.lib.json",
   "references": [
     {
-      "path": "../discojs"
+      "path": "../discojs/tsconfig.lib.json"
     }
   ],
   "compilerOptions": {
     "lib": ["DOM"],
     "rootDir": "./src",
-    "outDir": "dist"
+    "outDir": "dist",
+    "composite": true
   },
   "include": ["src"],
   "exclude": ["**/*.spec.ts"]

discojs-web/tsconfig.vitest.json

@@ -2,12 +2,13 @@
   "extends": "../tsconfig.base.json",
   "references": [
     {
-      "path": "../discojs"
+      "path": "../discojs/tsconfig.vitest.json"
     }
   ],
   "compilerOptions": {
     "lib": ["DOM"],
-    "noEmit": true
+    "noEmit": true,
+    "composite": true
   },
   "include": ["src"]
 }

discojs/package.json

@@ -2,11 +2,17 @@
   "name": "@epfml/discojs",
   "version": "3.0.0",
   "type": "module",
-  "exports": "./dist/index.js",
+  "exports": {
+    ".": {
+      "@disco/source": "./src/index.ts",
+      "default": "./dist/index.js",
+      "types": "./dist/index.d.ts"
+    }
+  },
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --exec pnpm run",
-    "build": "tsc --build",
+    "build": "tsc --build tsconfig.lib.json",
     "test": "cd .. && vitest --run --project=discojs"
   },
   "repository": {
@@ -16,6 +22,33 @@
   "bugs": {
     "url": "https://github.com/epfml/disco/issues"
   },
+  "imports": {
+    "#models/*": {
+      "@disco/source": "./src/models/*.ts",
+      "default": "./dist/models/*.js",
+      "types": "./dist/models/*.d.ts"
+    },
+    "#task/*": {
+      "@disco/source": "./src/task/*.ts",
+      "default": "./dist/task/*.js",
+      "types": "./dist/task/*.d.ts"
+    },
+    "#weights/*": {
+      "@disco/source": "./src/weights/*.ts",
+      "default": "./dist/weights/*.js",
+      "types": "./dist/weights/*.d.ts"
+    },
+    "#dtypes/*": {
+      "@disco/source": "./src/types/*.ts",
+      "default": "./dist/types/*.js",
+      "types": "./dist/types/*.d.ts"
+    },
+    "#dataset/*": {
+      "@disco/source": "./src/dataset/*.ts",
+      "default": "./dist/dataset/*.js",
+      "types": "./dist/dataset/*.d.ts"
+    }
+  },
   "homepage": "https://github.com/epfml/disco#readme",
   "dependencies": {
     "@epfml/isomorphic-wrtc": "workspace:*",

discojs/src/default_tasks/cifar10.ts

@@ -1,9 +1,5 @@
-import * as tf from "@tensorflow/tfjs";
-
-import type { Model, TaskProvider } from "../index.js";
-import { models } from "../index.js";
-
-import baseModel from "../models/mobileNet_v1_025_224.js";
+import type { TaskProvider } from "#task/index";
+import { cards } from "#models/index";
 
 export const cifar10: TaskProvider<"image", "decentralized"> = {
   getTask() {
@@ -67,28 +63,5 @@ export const cifar10: TaskProvider<"image", "decentralized"> = {
     });
   },
 
-  async getModel(): Promise<Model<"image">> {
-    const mobilenet = await tf.loadLayersModel({
-      load: async () => Promise.resolve(baseModel),
-    });
-
-    const x = mobilenet.getLayer("global_average_pooling2d_1");
-    const predictions = tf.layers
-      .dense({ units: 10, activation: "softmax", name: "denseModified" })
-      .apply(x.output) as tf.SymbolicTensor;
-
-    const model = tf.model({
-      inputs: mobilenet.input,
-      outputs: predictions,
-      name: "modelModified",
-    });
-
-    model.compile({
-      optimizer: "sgd",
-      loss: "categoricalCrossentropy",
-      metrics: ["accuracy"],
-    });
-
-    return new models.TFJS("image", model);
-  },
+  modelCard: cards.CIFAR10Classifier,
 };

discojs/src/default_tasks/lus_covid.ts

@@ -1,7 +1,5 @@
-import * as tf from "@tensorflow/tfjs";
-
-import type { Model, TaskProvider } from "../index.js";
-import { models } from "../index.js";
+import type { TaskProvider } from "#task/index";
+import { cards } from "#models/index";
 
 export const lusCovid: TaskProvider<"image", "federated"> = {
   getTask() {
@@ -46,65 +44,5 @@ export const lusCovid: TaskProvider<"image", "federated"> = {
 
   // Model architecture from tensorflow.js docs:
   // https://codelabs.developers.google.com/codelabs/tfjs-training-classfication/index.html#4
-  async getModel(): Promise<Model<"image">> {
-    const imageHeight = 100;
-    const imageWidth = 100;
-    const imageChannels = 3;
-    const numOutputClasses = 2;
-    const model = tf.sequential();
-
-    // In the first layer of our convolutional neural network we have
-    // to specify the input shape. Then we specify some parameters for
-    // the convolution operation that takes place in this layer.
-    model.add(
-      tf.layers.conv2d({
-        inputShape: [imageHeight, imageWidth, imageChannels],
-        kernelSize: 5,
-        filters: 8,
-        strides: 1,
-        activation: "relu",
-        kernelInitializer: "varianceScaling",
-      }),
-    );
-
-    // The MaxPooling layer acts as a sort of downsampling using max values
-    // in a region instead of averaging.
-    model.add(tf.layers.maxPooling2d({ poolSize: [2, 2], strides: [2, 2] }));
-
-    // Repeat the conv2d + maxPooling block.
-    // Note that we have more filters in the convolution.
-    model.add(
-      tf.layers.conv2d({
-        kernelSize: 5,
-        filters: 16,
-        strides: 1,
-        activation: "relu",
-        kernelInitializer: "varianceScaling",
-      }),
-    );
-    model.add(tf.layers.maxPooling2d({ poolSize: [2, 2], strides: [2, 2] }));
-
-    // Now we flatten the output from the 2D filters into a 1D vector to prepare
-    // it for input into our last layer. This is common practice when feeding
-    // higher dimensional data to a final classification output layer.
-    model.add(tf.layers.flatten());
-
-    // Our last layer is a dense layer which has 2 output units, one for each
-    // output class.
-    model.add(
-      tf.layers.dense({
-        units: numOutputClasses,
-        kernelInitializer: "varianceScaling",
-        activation: "softmax",
-      }),
-    );
-
-    model.compile({
-      optimizer: "sgd",
-      loss: "binaryCrossentropy",
-      metrics: ["accuracy"],
-    });
-
-    return Promise.resolve(new models.TFJS("image", model));
-  },
+  modelCard: cards.LUSClassifier,
 };

discojs/src/default_tasks/mnist.ts

@@ -1,7 +1,5 @@
-import * as tf from "@tensorflow/tfjs";
-
-import type { Model, TaskProvider } from "../index.js";
-import { models } from "../index.js";
+import type { TaskProvider } from "#task/index";
+import { cards } from "#models/index";
 
 export const mnist: TaskProvider<"image", "decentralized"> = {
   getTask() {
@@ -52,36 +50,5 @@ export const mnist: TaskProvider<"image", "decentralized"> = {
     });
   },
 
-  getModel(): Promise<Model<"image">> {
-    // Architecture from the PyTorch MNIST example (I made it slightly smaller, 650kB instead of 5MB)
-    // https://github.com/pytorch/examples/blob/main/mnist/main.py
-    const model = tf.sequential();
-
-    model.add(
-      tf.layers.conv2d({
-        inputShape: [28, 28, 3],
-        kernelSize: 5,
-        filters: 8,
-        activation: "relu",
-      }),
-    );
-    model.add(
-      tf.layers.conv2d({ kernelSize: 5, filters: 16, activation: "relu" }),
-    );
-    model.add(tf.layers.maxPooling2d({ poolSize: 2, strides: 2 }));
-    model.add(tf.layers.dropout({ rate: 0.25 }));
-
-    model.add(tf.layers.flatten());
-    model.add(tf.layers.dense({ units: 32, activation: "relu" }));
-    model.add(tf.layers.dropout({ rate: 0.25 }));
-    model.add(tf.layers.dense({ units: 10, activation: "softmax" }));
-
-    model.compile({
-      optimizer: "adam",
-      loss: "categoricalCrossentropy",
-      metrics: ["accuracy"],
-    });
-
-    return Promise.resolve(new models.TFJS("image", model));
-  },
+  modelCard: cards.PytorchMNISTClassifier,
 };