@epfml/discojs 3.0.1-p20250814105822.0 → 3.0.1-p20250924150135.0

package/dist/task/task.js

@@ -1,22 +1,50 @@
-import { isDisplayInformation } from './display_information.js';
-import { isTrainingInformation } from './training_information.js';
-export function isTaskID(obj) {
-    return typeof obj === 'string';
-}
-export function isTask(raw) {
-    if (typeof raw !== 'object' || raw === null) {
-        return false;
-    }
-    const { id, displayInformation, trainingInformation } = raw;
-    if (!isTaskID(id) ||
-        !isDisplayInformation(displayInformation) ||
-        !isTrainingInformation(trainingInformation)) {
-        return false;
-    }
-    const _ = {
-        id,
-        displayInformation,
-        trainingInformation,
+import { z } from "zod";
+import { DisplayInformation } from "./display_information.js";
+import { TrainingInformation } from "./training_information.js";
+export var Task;
+(function (Task) {
+    Task.baseSchema = z.object({
+        id: z.string(),
+        displayInformation: DisplayInformation.baseSchema,
+        trainingInformation: TrainingInformation.baseSchema,
+    });
+    Task.dataTypeToSchema = {
+        image: z.object({
+            dataType: z.literal("image"),
+            displayInformation: DisplayInformation.dataTypeToSchema.image,
+            trainingInformation: TrainingInformation.dataTypeToSchema.image,
+        }),
+        tabular: z.object({
+            dataType: z.literal("tabular"),
+            displayInformation: DisplayInformation.dataTypeToSchema.tabular,
+            trainingInformation: TrainingInformation.dataTypeToSchema.tabular,
+        }),
+        text: z.object({
+            dataType: z.literal("text"),
+            displayInformation: DisplayInformation.dataTypeToSchema.text,
+            trainingInformation: TrainingInformation.dataTypeToSchema.text,
+        }),
     };
-    return true;
-}
+    Task.networkToSchema = {
+        decentralized: z.object({
+            trainingInformation: TrainingInformation.networkToSchema.decentralized,
+        }),
+        federated: z.object({
+            trainingInformation: TrainingInformation.networkToSchema.federated,
+        }),
+        local: z.object({
+            trainingInformation: TrainingInformation.networkToSchema.local,
+        }),
+    };
+    Task.schema = Task.baseSchema
+        .and(z.union([
+        Task.dataTypeToSchema.image,
+        Task.dataTypeToSchema.tabular,
+        Task.dataTypeToSchema.text,
+    ]))
+        .and(z.union([
+        Task.networkToSchema.decentralized,
+        Task.networkToSchema.federated,
+        Task.networkToSchema.local,
+    ]));
+})(Task || (Task = {}));

package/dist/task/task_handler.d.ts

@@ -1,5 +1,5 @@
 import { Map } from "immutable";
-import type { DataType, Model } from "../index.js";
-import type { Task, TaskID } from "./task.js";
-export declare function pushTask<D extends DataType>(base: URL, task: Task<D>, model: Model<D>): Promise<void>;
-export declare function fetchTasks(base: URL): Promise<Map<TaskID, Task<DataType>>>;
+import type { DataType, Model, Network } from "../index.js";
+import type { Task } from "./task.js";
+export declare function pushTask<D extends DataType>(base: URL, task: Task<D, Network>, model: Model<D>): Promise<void>;
+export declare function fetchTasks(base: URL): Promise<Map<Task.ID, Task<DataType, Network>>>;

package/dist/task/task_handler.js

@@ -1,8 +1,5 @@
-import createDebug from "debug";
-import { Map } from "immutable";
+import { Map, Seq } from "immutable";
 import { serialization } from "../index.js";
-import { isTask } from "./task.js";
-const debug = createDebug("discojs:task:handlers");
 function urlToTasks(base) {
     const ret = new URL(base);
     ret.pathname += "tasks";
@@ -11,10 +8,10 @@ function urlToTasks(base) {
 export async function pushTask(base, task, model) {
     const response = await fetch(urlToTasks(base), {
         method: "POST",
+        headers: { "Content-Type": "application/json" },
         body: JSON.stringify({
-            task,
-            model: await serialization.model.encode(model),
-            weights: await serialization.weights.encode(model.weights),
+            task: serialization.task.serializeToJSON(task),
+            model: [...(await serialization.model.encode(model))],
         }),
     });
     if (!response.ok)
@@ -24,17 +21,16 @@ export async function fetchTasks(base) {
     const response = await fetch(urlToTasks(base));
     if (!response.ok)
         throw new Error(`fetch: HTTP status ${response.status}`);
-    const tasks = await response.json();
-    if (!Array.isArray(tasks)) {
-        throw new Error("Expected to receive an array of Tasks when fetching tasks");
+    const json = (await response.json());
+    if (!Array.isArray(json))
+        throw new Error("invalid tasks response: expected a JSON array");
+    const arr = json;
+    try {
+        return Map(Seq(await Promise.all(arr.map((t) => serialization.task.deserializeFromJSON(t)))).map((t) => [t.id, t]));
     }
-    else if (!tasks.every(isTask)) {
-        for (const task of tasks) {
-            if (!isTask(task)) {
-                debug("task has invalid format: :O", task);
-            }
-        }
-        throw new Error("invalid tasks response, the task object received is not well formatted");
+    catch (cause) {
+        throw new Error("invalid tasks response: unable to parse all tasks", {
+            cause,
+        });
     }
-    return Map(tasks.map((t) => [t.id, t]));
 }

package/dist/task/task_provider.d.ts

@@ -1,5 +1,5 @@
-import type { DataType, Model, Task } from "../index.js";
-export interface TaskProvider<D extends DataType> {
-    getTask(): Task<D>;
+import type { DataType, Model, Network, Task } from "../index.js";
+export interface TaskProvider<D extends DataType, N extends Network> {
+    getTask(): Promise<Task<D, N>>;
     getModel(): Promise<Model<D>>;
 }

package/dist/task/training_information.d.ts

@@ -1,38 +1,160 @@
-import { PreTrainedTokenizer } from "@xenova/transformers";
-import { DataType } from "../index.js";
-interface Privacy {
-    clippingRadius?: number;
-    noiseScale?: number;
-}
-export type TrainingInformation<D extends DataType> = {
-    epochs: number;
-    roundDuration: number;
-    validationSplit: number;
-    batchSize: number;
-    scheme: "decentralized" | "federated" | "local";
-    privacy?: Privacy;
-    maxShareValue?: number;
-    minNbOfParticipants: number;
-    aggregationStrategy?: "mean" | "secure";
-    tensorBackend: "tfjs" | "gpt";
-} & DataTypeToTrainingInformation[D];
-interface DataTypeToTrainingInformation {
-    image: {
-        dataType: "image";
-        LABEL_LIST: string[];
-        IMAGE_H: number;
-        IMAGE_W: number;
-    };
-    tabular: {
-        dataType: "tabular";
-        inputColumns: string[];
-        outputColumn: string;
+import { z } from "zod";
+import type { DataType, Network } from "../index.js";
+import { Tokenizer } from "../index.js";
+export declare namespace TrainingInformation {
+    const baseSchema: z.ZodObject<{
+        epochs: z.ZodNumber;
+        roundDuration: z.ZodNumber;
+        validationSplit: z.ZodNumber;
+        batchSize: z.ZodNumber;
+        tensorBackend: z.ZodEnum<["gpt", "tfjs"]>;
+    }, "strip", z.ZodTypeAny, {
+        epochs: number;
+        roundDuration: number;
+        validationSplit: number;
+        batchSize: number;
+        tensorBackend: "gpt" | "tfjs";
+    }, {
+        epochs: number;
+        roundDuration: number;
+        validationSplit: number;
+        batchSize: number;
+        tensorBackend: "gpt" | "tfjs";
+    }>;
+    const dataTypeToSchema: {
+        image: z.ZodObject<{
+            LABEL_LIST: z.ZodArray<z.ZodString, "many">;
+            IMAGE_W: z.ZodNumber;
+            IMAGE_H: z.ZodNumber;
+        }, "strip", z.ZodTypeAny, {
+            LABEL_LIST: string[];
+            IMAGE_W: number;
+            IMAGE_H: number;
+        }, {
+            LABEL_LIST: string[];
+            IMAGE_W: number;
+            IMAGE_H: number;
+        }>;
+        tabular: z.ZodObject<{
+            inputColumns: z.ZodArray<z.ZodString, "many">;
+            outputColumn: z.ZodString;
+        }, "strip", z.ZodTypeAny, {
+            inputColumns: string[];
+            outputColumn: string;
+        }, {
+            inputColumns: string[];
+            outputColumn: string;
+        }>;
+        text: z.ZodObject<{
+            tokenizer: z.ZodType<Tokenizer, z.ZodTypeDef, Tokenizer>;
+            contextLength: z.ZodNumber;
+        }, "strip", z.ZodTypeAny, {
+            contextLength: number;
+            tokenizer: Tokenizer;
+        }, {
+            contextLength: number;
+            tokenizer: Tokenizer;
+        }>;
     };
-    text: {
-        dataType: "text";
-        tokenizer: string | PreTrainedTokenizer;
-        contextLength: number;
+    const networkToSchema: {
+        decentralized: z.ZodIntersection<z.ZodObject<{
+            scheme: z.ZodLiteral<"decentralized">;
+        } & {
+            privacy: z.ZodOptional<z.ZodEffects<z.ZodObject<{
+                clippingRadius: z.ZodOptional<z.ZodNumber>;
+                noiseScale: z.ZodOptional<z.ZodNumber>;
+            }, "strip", z.ZodTypeAny, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            }, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            }>, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            } | undefined, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            }>>;
+            minNbOfParticipants: z.ZodNumber;
+        }, "strip", z.ZodTypeAny, {
+            scheme: "decentralized";
+            minNbOfParticipants: number;
+            privacy?: {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            } | undefined;
+        }, {
+            scheme: "decentralized";
+            minNbOfParticipants: number;
+            privacy?: {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            } | undefined;
+        }>, z.ZodUnion<[z.ZodObject<{
+            aggregationStrategy: z.ZodLiteral<"mean">;
+        }, "strip", z.ZodTypeAny, {
+            aggregationStrategy: "mean";
+        }, {
+            aggregationStrategy: "mean";
+        }>, z.ZodObject<{
+            aggregationStrategy: z.ZodLiteral<"secure">;
+            maxShareValue: z.ZodDefault<z.ZodOptional<z.ZodNumber>>;
+        }, "strip", z.ZodTypeAny, {
+            aggregationStrategy: "secure";
+            maxShareValue: number;
+        }, {
+            aggregationStrategy: "secure";
+            maxShareValue?: number | undefined;
+        }>]>>;
+        federated: z.ZodObject<{
+            scheme: z.ZodLiteral<"federated">;
+            aggregationStrategy: z.ZodLiteral<"mean">;
+        } & {
+            privacy: z.ZodOptional<z.ZodEffects<z.ZodObject<{
+                clippingRadius: z.ZodOptional<z.ZodNumber>;
+                noiseScale: z.ZodOptional<z.ZodNumber>;
+            }, "strip", z.ZodTypeAny, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            }, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            }>, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            } | undefined, {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            }>>;
+            minNbOfParticipants: z.ZodNumber;
+        }, "strip", z.ZodTypeAny, {
+            scheme: "federated";
+            minNbOfParticipants: number;
+            aggregationStrategy: "mean";
+            privacy?: {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            } | undefined;
+        }, {
+            scheme: "federated";
+            minNbOfParticipants: number;
+            aggregationStrategy: "mean";
+            privacy?: {
+                clippingRadius?: number | undefined;
+                noiseScale?: number | undefined;
+            } | undefined;
+        }>;
+        local: z.ZodObject<{
+            scheme: z.ZodLiteral<"local">;
+            aggregationStrategy: z.ZodLiteral<"mean">;
+        }, "strip", z.ZodTypeAny, {
+            scheme: "local";
+            aggregationStrategy: "mean";
+        }, {
+            scheme: "local";
+            aggregationStrategy: "mean";
+        }>;
     };
 }
-export declare function isTrainingInformation(raw: unknown): raw is TrainingInformation<DataType>;
-export {};
+export type TrainingInformation<D extends DataType, N extends Network> = z.infer<typeof TrainingInformation.baseSchema> & z.infer<(typeof TrainingInformation.dataTypeToSchema)[D]> & z.infer<(typeof TrainingInformation.networkToSchema)[N]>;

package/dist/task/training_information.js

@@ -1,112 +1,87 @@
-import { PreTrainedTokenizer } from "@xenova/transformers";
-function isPrivacy(raw) {
-    if (typeof raw !== "object" || raw === null) {
-        return false;
-    }
-    const { clippingRadius, noiseScale, } = raw;
-    if ((clippingRadius !== undefined && typeof clippingRadius !== "number") ||
-        (noiseScale !== undefined && typeof noiseScale !== "number"))
-        return false;
-    const _ = {
-        clippingRadius,
-        noiseScale,
+import { z } from "zod";
+import { Tokenizer } from "../index.js";
+const nonLocalNetworkSchema = z.object({
+    // reduce training accuracy and improve privacy.
+    privacy: z
+        .object({
+        // maximum weights difference between each round
+        clippingRadius: z.number().optional(),
+        // variance of the Gaussian noise added to the shared weights.
+        noiseScale: z.number().optional(),
+    })
+        .transform((o) => o.clippingRadius === undefined && o.noiseScale === undefined
+        ? undefined
+        : o)
+        .optional(),
+    // minimum number of participants required to train collaboratively
+    // In decentralized Learning the default is 3, in federated learning it is 2
+    minNbOfParticipants: z.number().positive().int(),
+});
+export var TrainingInformation;
+(function (TrainingInformation) {
+    TrainingInformation.baseSchema = z.object({
+        // number of epochs to run training for
+        epochs: z.number().positive().int(),
+        // number of epochs between each weight sharing round.
+        // e.g.if 3 then weights are shared every 3 epochs (in the distributed setting).
+        roundDuration: z.number().positive().int(),
+        // fraction of data to keep for validation, note this only works for image data
+        validationSplit: z.number().min(0).max(1),
+        // batch size of training data
+        batchSize: z.number().positive().int(),
+        // Tensor framework used by the model
+        tensorBackend: z.enum(["gpt", "tfjs"]),
+    });
+    TrainingInformation.dataTypeToSchema = {
+        image: z.object({
+            // classes, e.g. if two class of images, one with dogs and one with cats, then we would
+            // define ['dogs', 'cats'].
+            LABEL_LIST: z.array(z.string()).min(1),
+            // height of image to resize to
+            IMAGE_W: z.number().positive().int(),
+            // width of image to resize to
+            IMAGE_H: z.number().positive().int(),
+        }),
+        tabular: z.object({
+            // the columns to be chosen as input data for the model
+            inputColumns: z.array(z.string()),
+            // the columns to be predicted by the model
+            outputColumn: z.string(),
+        }),
+        text: z.object({
+            // should be set with the name of a Transformers.js pre-trained tokenizer, e.g., 'Xenova/gpt2'.
+            tokenizer: z.instanceof(Tokenizer),
+            // the maximum length of a input string used as input to a GPT model. It is used during preprocessing to
+            // truncate strings to a maximum length. The default value is tokenizer.model_max_length
+            contextLength: z.number().positive().int(),
+        }),
     };
-    return true;
-}
-export function isTrainingInformation(raw) {
-    if (typeof raw !== "object" || raw === null) {
-        return false;
-    }
-    const { aggregationStrategy, batchSize, dataType, privacy, epochs, maxShareValue, minNbOfParticipants, roundDuration, scheme, validationSplit, tensorBackend, } = raw;
-    if (typeof epochs !== "number" ||
-        typeof batchSize !== "number" ||
-        typeof roundDuration !== "number" ||
-        typeof validationSplit !== "number" ||
-        typeof minNbOfParticipants !== "number" ||
-        (privacy !== undefined && !isPrivacy(privacy)) ||
-        (maxShareValue !== undefined && typeof maxShareValue !== "number")) {
-        return false;
-    }
-    switch (aggregationStrategy) {
-        case undefined:
-        case "mean":
-        case "secure":
-            break;
-        default:
-            return false;
-    }
-    switch (tensorBackend) {
-        case "tfjs":
-        case "gpt":
-            break;
-        default:
-            return false;
-    }
-    switch (scheme) {
-        case "decentralized":
-        case "federated":
-        case "local":
-            break;
-        default:
-            return false;
-    }
-    const repack = {
-        aggregationStrategy,
-        batchSize,
-        epochs,
-        maxShareValue,
-        minNbOfParticipants,
-        privacy,
-        roundDuration,
-        scheme,
-        tensorBackend,
-        validationSplit,
+    TrainingInformation.networkToSchema = {
+        decentralized: z
+            .object({
+            scheme: z.literal("decentralized"),
+        })
+            .merge(nonLocalNetworkSchema)
+            .and(z.union([
+            z.object({
+                aggregationStrategy: z.literal("mean"),
+            }),
+            z.object({
+                aggregationStrategy: z.literal("secure"),
+                // Secure Aggregation: maximum absolute value of a number in a randomly generated share
+                // default is 100, must be a positive number, check the docs/PRIVACY.md file for more information on significance of maxShareValue selection
+                maxShareValue: z.number().positive().int().optional().default(100),
+            }),
+        ])),
+        federated: z
+            .object({
+            scheme: z.literal("federated"),
+            aggregationStrategy: z.literal("mean"),
+        })
+            .merge(nonLocalNetworkSchema),
+        local: z.object({
+            scheme: z.literal("local"),
+            aggregationStrategy: z.literal("mean"),
+        }),
     };
-    switch (dataType) {
-        case "image": {
-            const { LABEL_LIST, IMAGE_W, IMAGE_H } = raw;
-            if (!(Array.isArray(LABEL_LIST) &&
-                LABEL_LIST.every((e) => typeof e === "string")) ||
-                typeof IMAGE_H !== "number" ||
-                typeof IMAGE_W !== "number")
-                return false;
-            const _ = {
-                ...repack,
-                dataType,
-                LABEL_LIST,
-                IMAGE_W,
-                IMAGE_H,
-            };
-            return true;
-        }
-        case "tabular": {
-            const { inputColumns, outputColumn } = raw;
-            if (!(Array.isArray(inputColumns) &&
-                inputColumns.every((e) => typeof e === "string")) ||
-                typeof outputColumn !== "string")
-                return false;
-            const _ = {
-                ...repack,
-                dataType,
-                inputColumns,
-                outputColumn,
-            };
-            return true;
-        }
-        case "text": {
-            const { contextLength, tokenizer, } = raw;
-            if ((typeof tokenizer !== "string" &&
-                !(tokenizer instanceof PreTrainedTokenizer)) ||
-                (typeof contextLength !== "number"))
-                return false;
-            const _ = {
-                ...repack,
-                dataType,
-                contextLength,
-                tokenizer,
-            };
-            return true;
-        }
-    }
-    return false;
-}
+})(TrainingInformation || (TrainingInformation = {}));

package/dist/training/disco.d.ts

@@ -1,10 +1,10 @@
-import { client as clients, BatchLogs, EpochLogs, Logger, TrainingInformation, Dataset } from "../index.js";
-import type { DataFormat, DataType, Task } from "../index.js";
+import { client as clients, BatchLogs, EpochLogs, Logger, Dataset } from "../index.js";
+import type { DataFormat, DataType, Network, Task } from "../index.js";
 import type { Aggregator } from "../aggregator/index.js";
 import { EventEmitter } from "../utils/event_emitter.js";
 import { RoundLogs, Trainer } from "./trainer.js";
-interface DiscoConfig {
-    scheme: TrainingInformation<DataType>["scheme"];
+interface DiscoConfig<N extends Network> {
+    scheme: N;
     logger: Logger;
     /**
      * keep preprocessed dataset in memory while training
@@ -24,12 +24,12 @@ export type RoundStatus = 'not enough participants' | // Server notification to
  * a convenient object providing a reduced yet complete API that wraps model training and
  * communication with nodes.
  */
-export declare class Disco<D extends DataType> extends EventEmitter<{
+export declare class Disco<D extends DataType, N extends Network> extends EventEmitter<{
     status: RoundStatus;
     participants: number;
 }> {
     #private;
-    readonly trainer: Trainer<D>;
+    readonly trainer: Trainer<D, N>;
     /**
      * Connect to the given task and get ready to train.
      *
@@ -37,10 +37,10 @@ export declare class Disco<D extends DataType> extends EventEmitter<{
      * @param clientConfig client to connect with or parameters on how to create one.
      * @param config the DiscoConfig
      */
-    constructor(task: Task<D>, clientConfig: clients.Client | URL | {
+    constructor(task: Task<D, N>, clientConfig: clients.Client<N> | URL | {
         aggregator: Aggregator;
         url: URL;
-    }, config: Partial<DiscoConfig>);
+    }, config: Partial<DiscoConfig<N>>);
     /** Train on dataset, yielding logs of every round. */
     trainByRound(dataset: Dataset<DataFormat.Raw[D]>): AsyncGenerator<RoundLogs>;
     /** Train on dataset, yielding logs of every epoch. */

package/dist/training/disco.js

@@ -24,6 +24,7 @@ export class Disco extends EventEmitter {
     constructor(task, clientConfig, config) {
         super();
         const { scheme, logger, preprocessOnce } = {
+            // cast as typescript is bad at generic
             scheme: task.trainingInformation.scheme,
             logger: new ConsoleLogger(),
             preprocessOnce: false,
@@ -135,7 +136,7 @@ export class Disco extends EventEmitter {
     }
     async #preprocessSplitAndBatch(dataset) {
         const { batchSize, validationSplit } = this.#task.trainingInformation;
-        let preprocessed = await processing.preprocess(this.#task, dataset);
+        let preprocessed = processing.preprocess(this.#task, dataset);
         preprocessed = (this.#preprocessOnce
             ? new Dataset(await arrayFromAsync(preprocessed))
             : preprocessed);

package/dist/training/trainer.d.ts

@@ -1,16 +1,16 @@
 import { List } from "immutable";
-import type { Batched, BatchLogs, Dataset, DataFormat, DataType, EpochLogs, Model, Task } from "../index.js";
+import type { Batched, BatchLogs, Dataset, DataFormat, DataType, EpochLogs, Model, Task, Network } from "../index.js";
 import { Client } from "../client/index.js";
 export interface RoundLogs {
     epochs: List<EpochLogs>;
     participants: number;
 }
 /** Train a model and exchange with others **/
-export declare class Trainer<D extends DataType> {
+export declare class Trainer<D extends DataType, N extends Network> {
     #private;
     get model(): Model<D>;
     set model(model: Model<D>);
-    constructor(task: Task<D>, client: Client);
+    constructor(task: Task<D, N>, client: Client<N>);
     stopTraining(): Promise<void>;
     train(dataset: Dataset<Batched<DataFormat.ModelEncoded[D]>>, validationDataset?: Dataset<Batched<DataFormat.ModelEncoded[D]>>): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs>, void>;
 }

package/dist/training/trainer.js

@@ -22,7 +22,8 @@ export class Trainer {
         this.#client = client;
         this.#roundDuration = task.trainingInformation.roundDuration;
         this.#epochs = task.trainingInformation.epochs;
-        this.#privacy = task.trainingInformation.privacy;
+        if ("privacy" in task.trainingInformation)
+            this.#privacy = task.trainingInformation.privacy;
         if (!Number.isInteger(this.#epochs / this.#roundDuration))
             throw new Error(`round duration ${this.#roundDuration} doesn't divide number of epochs ${this.#epochs}`);
     }

package/dist/types/index.d.ts

@@ -1,2 +1,3 @@
 export * as DataFormat from "./data_format.js";
 export type DataType = "image" | "tabular" | "text";
+export type Network = "decentralized" | "federated" | "local";

package/dist/validator.d.ts

@@ -1,10 +1,10 @@
-import { Dataset, DataFormat, DataType, Model, Task } from "./index.js";
+import type { Dataset, DataFormat, DataType, Model, Task, Network } from "./index.js";
 export declare class Validator<D extends DataType> {
     #private;
-    readonly task: Task<D>;
-    constructor(task: Task<D>, model: Model<D>);
+    readonly task: Task<D, Network>;
+    constructor(task: Task<D, Network>, model: Model<D>);
     /** infer every line of the dataset and check that it is as labelled */
-    test(dataset: Dataset<DataFormat.Raw[D]>): Promise<Dataset<Record<"predicted" | "truth", DataFormat.Inferred[D]>>>;
+    test(dataset: Dataset<DataFormat.Raw[D]>): Dataset<Record<"predicted" | "truth", DataFormat.Inferred[D]>>;
     /** use the model to predict every line of the dataset */
     infer(dataset: Dataset<DataFormat.RawWithoutLabel[D]>): AsyncGenerator<DataFormat.Inferred[D], void>;
 }