Break import cycles

.github/workflows/_build.yml

@@ -3,37 +3,6 @@ on:
   workflow_call:
 
 jobs:
-  build-lib:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs run build
-
-  build-lib-node:
-    needs: build-lib
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs-node run build
-
-  build-lib-web:
-    needs: build-lib
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs-web run build
-
-  build-server:
-    needs: build-lib-node
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F server run build
-
   build-server-docker:
     runs-on: ubuntu-latest
     steps:
@@ -51,29 +20,11 @@ jobs:
             timeout=$((timeout - 1))
           done
 
-  build-cli:
-    needs: build-server
+  build-topological:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup-node
-      # Build all subprojects to upload all artifacts at once
       - run: pnpm -r run build
       - uses: actions/upload-artifact@v7
         with: { name: all-builds, path: "*/dist" }
-
-  build-webapp:
-    needs: build-lib-web
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F webapp run build
-
-  build-docs-examples:
-    needs: build-server
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F examples run build

.github/workflows/_static-analysis.yml

@@ -16,3 +16,10 @@ jobs:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup-node
       - run: pnpm exec knip
+
+  check_cycles:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: ./.github/actions/setup-node
+      - run: pnpm run check_cycles

cli/src/benchmark_gpt.ts

@@ -2,13 +2,13 @@ import "@tensorflow/tfjs-node";
 import { List } from "immutable";
 import { parse } from "ts-command-line-args";
 
-import type { Network, Task } from "@epfml/discojs";
+import type { Network, Task, GPTConfig } from "@epfml/discojs";
 import {
-  async_iterator,
+  gather,
   defaultTasks,
   defaultModels,
   fetchTasks,
-  models,
+  GPT,
 } from "@epfml/discojs";
 import { loadModelFromDisk, loadText } from "@epfml/discojs-node";
 
@@ -109,8 +109,8 @@ async function main(args: Required<CLIArguments>): Promise<void> {
     const epochsCount = 1;
     const iterationsPerEpoch = 10;
 
-    const config: models.GPTConfig = {
-      modelType: modelType as models.GPTConfig["modelType"],
+    const config: GPTConfig = {
+      modelType: modelType as GPTConfig["modelType"],
       maxIter: iterationsPerEpoch,
       lr: 0.0001,
       contextLength,
@@ -130,16 +130,14 @@ async function main(args: Required<CLIArguments>): Promise<void> {
       .batch(batchSize);
 
     // Init and train the model
-    const model = new models.GPT(config);
+    const model = new GPT(config);
     console.log(
       `\tmodel type ${modelType} \n\tbatch size ${batchSize} \n\tcontext length ${contextLength}`,
     );
 
     let epochTime = performance.now();
     for (let epochsCounter = 1; epochsCounter <= epochsCount; epochsCounter++) {
-      const [_, logs] = await async_iterator.gather(
-        model.train(preprocessedDataset),
-      );
+      const [_, logs] = await gather(model.train(preprocessedDataset));
       epochTime = performance.now() - epochTime;
       const msPerToken =
         epochTime /
@@ -154,7 +152,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
      */
   } else {
     const model = await loadModelFromDisk(modelPath);
-    if (!(model instanceof models.GPT)) {
+    if (!(model instanceof GPT)) {
       throw new Error("Loaded model isn't a GPT model");
     }
 

cli/src/cli.ts

@@ -15,11 +15,7 @@ import type {
   TaskProvider,
   Network,
 } from "@epfml/discojs";
-import {
-  Disco,
-  aggregator as aggregators,
-  client as clients,
-} from "@epfml/discojs";
+import { Disco, getAggregator, getClient } from "@epfml/discojs";
 
 import { getTaskData } from "./data.js";
 import { args } from "./args.js";
@@ -35,8 +31,8 @@ async function runUser<D extends DataType, N extends Network>(
 ): Promise<List<SummaryLogs>> {
   // cast as typescript isn't good with generics
   const trainingScheme = task.trainingInformation.scheme as N;
-  const aggregator = aggregators.getAggregator(task);
-  const client = clients.getClient(trainingScheme, url, task, aggregator);
+  const aggregator = getAggregator(task);
+  const client = getClient(trainingScheme, url, task, aggregator);
   const disco = new Disco(task, client, { scheme: trainingScheme });
 
   const dir = path.join(".", `${args.testID}`);

cli/src/data.ts

@@ -1,6 +1,7 @@
 import path from "node:path";
-import { Dataset, processing } from "@epfml/discojs";
-import { DataFormat, DataType, Image, Task } from "@epfml/discojs";
+import type { Dataset } from "@epfml/discojs";
+import { extractColumn } from "@epfml/discojs";
+import type { DataFormat, DataType, Image, Task } from "@epfml/discojs";
 import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node";
 import { Repeat } from "immutable";
 
@@ -31,10 +32,7 @@ function loadTinderDogData(split: number): Dataset<DataFormat.Raw["image"]> {
   return loadCSV(path.join(folder, "labels.csv"))
     .map(
       (row) =>
-        [
-          processing.extractColumn(row, "filename"),
-          processing.extractColumn(row, "label"),
-        ] as const,
+        [extractColumn(row, "filename"), extractColumn(row, "label")] as const,
     )
     .map(async ([filename, label]) => {
       try {
@@ -58,10 +56,7 @@ function loadData(
   return loadCSV(path.join(folder, "labels.csv"))
     .map(
       (row) =>
-        [
-          processing.extractColumn(row, "filename"),
-          processing.extractColumn(row, "label"),
-        ] as const,
+        [extractColumn(row, "filename"), extractColumn(row, "label")] as const,
     )
     .map(async ([filename, label]) => {
       try {

cli/src/hellaswag_gpt.ts

@@ -5,22 +5,25 @@ import { parse } from "ts-command-line-args";
 
 import "@tensorflow/tfjs-node";
 import path from "node:path";
-import { models, serialization, Tokenizer } from "@epfml/discojs";
+import type { HellaSwagDataset } from "@epfml/discojs";
+import {
+  GPT,
+  ONNXModel,
+  modelDecode,
+  Tokenizer,
+  evaluate_hellaswag,
+} from "@epfml/discojs";
 import { loadHellaSwag } from "@epfml/discojs-node";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 
-async function evaluateModel(
-  model: models.GPT | models.ONNXModel,
-  numDataPoints = -1,
-) {
-  const hellaswagDataset: models.HellaSwagDataset =
-    await loadHellaSwag(numDataPoints);
+async function evaluateModel(model: GPT | ONNXModel, numDataPoints = -1) {
+  const hellaswagDataset: HellaSwagDataset = await loadHellaSwag(numDataPoints);
   const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2");
   console.log("Starting the HellaSwag benchmark...");
 
   const start = Date.now();
-  const accuracy = await models.evaluate_hellaswag(
+  const accuracy = await evaluate_hellaswag(
     model,
     tokenizer,
     hellaswagDataset,
@@ -91,15 +94,15 @@ async function main(): Promise<void> {
     { helpArg: "help" },
   );
 
-  let model: models.GPT | models.ONNXModel | undefined;
+  let model: GPT | ONNXModel | undefined;
   switch (args.model) {
     case "onnx":
       console.log("Using ONNX pretrained model Xenova/gpt2");
-      model = await models.ONNXModel.init_pretrained("Xenova/gpt2");
+      model = await ONNXModel.init_pretrained("Xenova/gpt2");
       break;
     case "gpt-tfjs-random":
       console.log("Using GPT-TFJS with random initialization");
-      model = new models.GPT({ seed: 42 });
+      model = new GPT({ seed: 42 });
       break;
     case "gpt-tfjs-pretrained":
       console.log("Using GPT-TFJS with pretrained weights");
@@ -109,7 +112,7 @@ async function main(): Promise<void> {
         );
       }
       const encodedModel = await fsPromise.readFile(args.pretrainedModelPath);
-      model = (await serialization.model.decode(encodedModel)) as models.GPT;
+      model = (await modelDecode(encodedModel)) as GPT;
       break;
   }
   await evaluateModel(model, args.numDataPoints);

cli/src/train_gpt.ts

@@ -1,12 +1,13 @@
 import "@tensorflow/tfjs-node";
-import { models, Dataset, Tokenizer } from "@epfml/discojs";
+import type { GPTConfig } from "@epfml/discojs";
+import { GPT, Dataset, Tokenizer } from "@epfml/discojs";
 import { List } from "immutable";
 
 async function main(): Promise<void> {
   const data = "Lorem ipsum dolor sit amet, consectetur adipis";
   const seed = 42;
 
-  const config: models.GPTConfig = {
+  const config: GPTConfig = {
     modelType: "gpt-nano",
     lr: 0.01,
     maxIter: 50,
@@ -26,7 +27,7 @@ async function main(): Promise<void> {
     .repeat()
     .batch(8);
 
-  const model = new models.GPT(config);
+  const model = new GPT(config);
   for await (const logs of model.train(tokenDataset, undefined)) {
     console.log(logs);
   }

cli/src/user_log.ts

@@ -1,4 +1,5 @@
-import { args, BenchmarkArguments } from "./args.js";
+import type { BenchmarkArguments } from "./args.js";
+import { args } from "./args.js";
 import type { SummaryLogs, DataType, Network, Task } from "@epfml/discojs";
 
 type SerializableArguments = Omit<BenchmarkArguments, "provider" | "host"> & {

discojs-node/package.json

@@ -10,7 +10,7 @@
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs/dist --watch . --exec pnpm run",
-    "build": "tsc --build tsconfig.lib.json",
+    "build": "tsc --build",
     "test": "cd .. && vitest --run --project=discojs-node"
   },
   "repository": {

discojs-node/src/hellaswag.ts

@@ -2,7 +2,8 @@ import path from "node:path";
 import fetch from "node-fetch";
 import fs from "node:fs/promises";
 
-import { models } from "@epfml/discojs";
+import type { HellaSwagExample, HellaSwagDataset } from "@epfml/discojs";
+import { HELLASWAG_URL } from "@epfml/discojs";
 
 import { dirname } from "path";
 import { fileURLToPath } from "url";
@@ -17,18 +18,18 @@ const hellaswag_filepath = path.join(DATASET_DIR, "hellaswag_val.jsonl");
  * @param limit - Maximum number of examples to load (-1 means all)
  * @returns A HellaSwagDataset containing the examples.
  */
-export async function load(limit = -1): Promise<models.HellaSwagDataset> {
+export async function load(limit = -1): Promise<HellaSwagDataset> {
   let text: string;
   try {
     // Reads the file if it exists locally
     text = (await fs.readFile(hellaswag_filepath)).toString();
   } catch {
     console.log("Downloading the Hellaswag benchmark");
     // Otherwise fetch it
-    const response = await fetch(models.HELLASWAG_URL);
+    const response = await fetch(HELLASWAG_URL);
     if (!response.ok) {
       throw new Error(
-        `Failed to fetch dataset from ${models.HELLASWAG_URL}: ${response.statusText}`,
+        `Failed to fetch dataset from ${HELLASWAG_URL}: ${response.statusText}`,
       );
     }
 
@@ -39,14 +40,14 @@ export async function load(limit = -1): Promise<models.HellaSwagDataset> {
 
   const lines = text.split("\n");
 
-  const dataset: models.HellaSwagDataset = [];
+  const dataset: HellaSwagDataset = [];
   let count = 0;
   for (const line of lines) {
     if (line.trim().length === 0) continue;
     if (limit !== -1 && count >= limit) break;
 
     try {
-      const data = JSON.parse(line.trim()) as models.HellaSwagExample;
+      const data = JSON.parse(line.trim()) as HellaSwagExample;
       dataset.push(data);
       count++;
     } catch (e) {

discojs-node/src/loaders/text.ts

@@ -1,6 +1,7 @@
 import createDebug from "debug";
 import { createReadStream } from "node:fs";
-import { Dataset, Text } from "@epfml/discojs";
+import type { Text } from "@epfml/discojs";
+import { Dataset } from "@epfml/discojs";
 
 const debug = createDebug("discojs-node:loaders:text");
 

discojs-node/src/model_loader.ts

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

discojs-node/tsconfig.json

@@ -1,7 +1,4 @@
 {
-  "compilerOptions": {
-    "composite": true
-  },
   "files": [],
   "references": [
     {