@epfml/discojs 2.1.1 → 2.1.2-p20240506085037.0

package/dist/{core/privacy.js → privacy.js}

@@ -1,7 +1,4 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.addDifferentialPrivacy = void 0;
-var _1 = require(".");
+import * as tf from '@tensorflow/tfjs';
 /**
  * Add task-parametrized Gaussian noise to and clip the weights update between the previous and current rounds.
  * The previous round's weights are the last weights pulled from server/peers.
@@ -11,20 +8,19 @@ var _1 = require(".");
  * @param task the task
  * @returns the noised weights for the current round
  */
-function addDifferentialPrivacy(updatedWeights, staleWeights, task) {
-    var _a, _b;
-    var noiseScale = (_a = task.trainingInformation) === null || _a === void 0 ? void 0 : _a.noiseScale;
-    var clippingRadius = (_b = task.trainingInformation) === null || _b === void 0 ? void 0 : _b.clippingRadius;
-    var weightsDiff = updatedWeights.sub(staleWeights);
-    var newWeightsDiff;
+export function addDifferentialPrivacy(updatedWeights, staleWeights, task) {
+    const noiseScale = task.trainingInformation?.noiseScale;
+    const clippingRadius = task.trainingInformation?.clippingRadius;
+    const weightsDiff = updatedWeights.sub(staleWeights);
+    let newWeightsDiff;
     if (clippingRadius !== undefined) {
         // Frobenius norm
-        var norm_1 = weightsDiff.frobeniusNorm();
-        newWeightsDiff = weightsDiff.map(function (w) {
-            var clipped = w.div(Math.max(1, norm_1 / clippingRadius));
+        const norm = weightsDiff.frobeniusNorm();
+        newWeightsDiff = weightsDiff.map((w) => {
+            const clipped = w.div(Math.max(1, norm / clippingRadius));
             if (noiseScale !== undefined) {
                 // Add clipping and noise
-                var noise = _1.tf.randomNormal(w.shape, 0, (noiseScale * noiseScale) * (clippingRadius * clippingRadius));
+                const noise = tf.randomNormal(w.shape, 0, (noiseScale * noiseScale) * (clippingRadius * clippingRadius));
                 return clipped.add(noise);
             }
             else {
@@ -36,7 +32,7 @@ function addDifferentialPrivacy(updatedWeights, staleWeights, task) {
     else {
         if (noiseScale !== undefined) {
             // Add noise without any clipping
-            newWeightsDiff = weightsDiff.map(function (w) { return _1.tf.randomNormal(w.shape, 0, (noiseScale * noiseScale)); });
+            newWeightsDiff = weightsDiff.map((w) => tf.randomNormal(w.shape, 0, (noiseScale * noiseScale)));
         }
         else {
             return updatedWeights;
@@ -44,4 +40,3 @@ function addDifferentialPrivacy(updatedWeights, staleWeights, task) {
     }
     return staleWeights.add(newWeightsDiff);
 }
-exports.addDifferentialPrivacy = addDifferentialPrivacy;

package/dist/serialization/index.d.ts

@@ -0,0 +1,2 @@
+export * as model from './model.js';
+export * as weights from './weights.js';

package/dist/serialization/index.js

@@ -0,0 +1,2 @@
+export * as model from './model.js';
+export * as weights from './weights.js';

package/dist/serialization/model.d.ts

@@ -0,0 +1,5 @@
+import type { Model } from '../index.js';
+export type Encoded = Uint8Array;
+export declare function isEncoded(raw: unknown): raw is Encoded;
+export declare function encode(model: Model): Promise<Encoded>;
+export declare function decode(encoded: unknown): Promise<Model>;

package/dist/serialization/model.js

@@ -0,0 +1,67 @@
+import msgpack from 'msgpack-lite';
+import { models, serialization } from '../index.js';
+const Type = {
+    TFJS: 0,
+    GPT: 1
+};
+export function isEncoded(raw) {
+    return raw instanceof Uint8Array;
+}
+export async function encode(model) {
+    if (model instanceof models.TFJS) {
+        const serialized = await model.serialize();
+        return msgpack.encode([Type.TFJS, serialized]);
+    }
+    if (model instanceof models.GPT) {
+        const { weights, config } = model.serialize();
+        const serializedWeights = await serialization.weights.encode(weights);
+        return msgpack.encode([Type.GPT, serializedWeights, config]);
+    }
+    throw new Error('unknown model type');
+}
+export async function decode(encoded) {
+    if (!isEncoded(encoded)) {
+        throw new Error("Invalid encoding, raw encoding isn't an instance of Uint8Array");
+    }
+    const raw = msgpack.decode(encoded);
+    if (!Array.isArray(raw) || raw.length < 2) {
+        throw new Error("invalid encoding, encoding isn't an array or doesn't contain enough values");
+    }
+    const type = raw[0];
+    if (typeof type !== 'number') {
+        throw new Error('invalid encoding, first encoding field should be the model type');
+    }
+    const rawModel = raw[1];
+    switch (type) {
+        case Type.TFJS:
+            if (raw.length !== 2) {
+                throw new Error('invalid encoding, TFJS model encoding should be an array of length 2');
+            }
+            // TODO totally unsafe casting
+            return await models.TFJS.deserialize(rawModel);
+        case Type.GPT: {
+            let config;
+            if (raw.length == 2) {
+                config = undefined;
+            }
+            else if (raw.length == 3) {
+                config = raw[2];
+            }
+            else {
+                throw new Error('invalid encoding, gpt-tfjs model encoding should be an array of length 2 or 3');
+            }
+            if (!Array.isArray(rawModel)) {
+                throw new Error('invalid encoding, gpt-tfjs model weights should be an array');
+            }
+            const arr = rawModel;
+            if (arr.some((r) => typeof r !== 'number')) {
+                throw new Error("invalid encoding, gpt-tfjs weights should be numbers");
+            }
+            const nums = arr;
+            const weights = serialization.weights.decode(nums);
+            return models.GPT.deserialize({ weights, config });
+        }
+        default:
+            throw new Error('invalid encoding, model type unrecognized');
+    }
+}

package/dist/{core/serialization → serialization}/weights.d.ts

@@ -1,5 +1,5 @@
-import { WeightsContainer } from '..';
-export declare type Encoded = number[];
+import { WeightsContainer } from '../index.js';
+export type Encoded = number[];
 export declare function isEncoded(raw: unknown): raw is Encoded;
 export declare function encode(weights: WeightsContainer): Promise<Encoded>;
 export declare function decode(encoded: Encoded): WeightsContainer;

package/dist/serialization/weights.js

@@ -0,0 +1,37 @@
+import * as msgpack from 'msgpack-lite';
+import * as tf from '@tensorflow/tfjs';
+import { WeightsContainer } from '../index.js';
+function isSerialized(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { shape, data } = raw;
+    if (!(Array.isArray(shape) && shape.every((e) => typeof e === 'number')) ||
+        !(Array.isArray(data) && data.every((e) => typeof e === 'number'))) {
+        return false;
+    }
+    const _ = {
+        shape: shape,
+        data: data,
+    };
+    return true;
+}
+export function isEncoded(raw) {
+    return Array.isArray(raw) && raw.every((e) => typeof e === 'number');
+}
+export async function encode(weights) {
+    const serialized = await Promise.all(weights.weights.map(async (t) => {
+        return {
+            shape: t.shape,
+            data: [...await t.data()]
+        };
+    }));
+    return [...msgpack.encode(serialized).values()];
+}
+export function decode(encoded) {
+    const raw = msgpack.decode(encoded);
+    if (!(Array.isArray(raw) && raw.every(isSerialized))) {
+        throw new Error('expected to decode an array of serialized weights');
+    }
+    return new WeightsContainer(raw.map((w) => tf.tensor(w.data, w.shape)));
+}

package/dist/task/data_example.js

@@ -0,0 +1,14 @@
+export function isDataExample(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { columnName, columnData } = raw;
+    if (typeof columnName !== 'string' ||
+        (typeof columnData !== 'string' && typeof columnData !== 'number')) {
+        return false;
+    }
+    const repack = { columnName, columnData };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/task/digest.js

@@ -0,0 +1,14 @@
+export function isDigest(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { algorithm, value } = raw;
+    if (!(typeof algorithm === 'string' &&
+        typeof value === 'string')) {
+        return false;
+    }
+    const repack = { algorithm, value };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/{core/task → task}/display_information.d.ts

@@ -1,6 +1,6 @@
-import { Summary } from './summary';
-import { DataExample } from './data_example';
-export declare function isDisplayInformation(raw: unknown): raw is DisplayInformation;
+import { type Summary } from './summary.js';
+import { type DataExample } from './data_example.js';
+import { type LabelType } from './label_type.js';
 export interface DisplayInformation {
     taskTitle?: string;
     summary?: Summary;
@@ -12,4 +12,6 @@ export interface DisplayInformation {
     headers?: string[];
     dataExampleImage?: string;
     limitations?: string;
+    labelDisplay?: LabelType;
 }
+export declare function isDisplayInformation(raw: unknown): raw is DisplayInformation;

package/dist/task/display_information.js

@@ -0,0 +1,46 @@
+import { isSummary } from './summary.js';
+import { isDataExample } from './data_example.js';
+import { isLabelType } from './label_type.js';
+export function isDisplayInformation(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { dataExample, dataExampleImage, dataExampleText, dataFormatInformation, headers, labelDisplay, limitations, model, summary, taskTitle, tradeoffs } = raw;
+    if (typeof taskTitle !== 'string' ||
+        (dataExampleText !== undefined && typeof dataExampleText !== 'string') ||
+        (dataFormatInformation !== undefined && typeof dataFormatInformation !== 'string') ||
+        (tradeoffs !== undefined && typeof tradeoffs !== 'string') ||
+        (model !== undefined && typeof model !== 'string') ||
+        (dataExampleImage !== undefined && typeof dataExampleImage !== 'string') ||
+        (labelDisplay !== undefined && !isLabelType(labelDisplay)) ||
+        (limitations !== undefined && typeof limitations !== 'string')) {
+        return false;
+    }
+    if (summary !== undefined && !isSummary(summary)) {
+        return false;
+    }
+    if (dataExample !== undefined && !(Array.isArray(dataExample) &&
+        dataExample.every(isDataExample))) {
+        return false;
+    }
+    if (headers !== undefined && !(Array.isArray(headers) &&
+        headers.every((e) => typeof e === 'string'))) {
+        return false;
+    }
+    const repack = {
+        dataExample,
+        dataExampleImage,
+        dataExampleText,
+        dataFormatInformation,
+        headers,
+        labelDisplay,
+        limitations,
+        model,
+        summary,
+        taskTitle,
+        tradeoffs,
+    };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/task/index.d.ts

@@ -0,0 +1,7 @@
+export { isTask, type Task, isTaskID, type TaskID } from './task.js';
+export { type TaskProvider } from './task_provider.js';
+export { isDigest, type Digest } from './digest.js';
+export { isDisplayInformation, type DisplayInformation } from './display_information.js';
+export type { TrainingInformation } from './training_information.js';
+export { pushTask, fetchTasks } from './task_handler.js';
+export { LabelTypeEnum } from './label_type.js';

package/dist/task/index.js

@@ -0,0 +1,5 @@
+export { isTask, isTaskID } from './task.js';
+export { isDigest } from './digest.js';
+export { isDisplayInformation } from './display_information.js';
+export { pushTask, fetchTasks } from './task_handler.js';
+export { LabelTypeEnum } from './label_type.js';

package/dist/task/label_type.d.ts

@@ -0,0 +1,9 @@
+export interface LabelType {
+    labelType: LabelTypeEnum;
+    mapBaseUrl?: string;
+}
+export declare enum LabelTypeEnum {
+    TEXT = 0,
+    POLYGON_MAP = 1
+}
+export declare function isLabelType(raw: unknown): raw is LabelType;

package/dist/task/label_type.js

@@ -0,0 +1,28 @@
+export var LabelTypeEnum;
+(function (LabelTypeEnum) {
+    LabelTypeEnum[LabelTypeEnum["TEXT"] = 0] = "TEXT";
+    LabelTypeEnum[LabelTypeEnum["POLYGON_MAP"] = 1] = "POLYGON_MAP";
+})(LabelTypeEnum || (LabelTypeEnum = {}));
+function isLabelTypeEnum(raw) {
+    switch (raw) {
+        case LabelTypeEnum.TEXT: break;
+        case LabelTypeEnum.POLYGON_MAP: break;
+        default: return false;
+    }
+    const _ = raw;
+    return true;
+}
+export function isLabelType(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { labelType, mapBaseUrl } = raw;
+    if (!isLabelTypeEnum(labelType) ||
+        (mapBaseUrl !== undefined && typeof mapBaseUrl !== 'string')) {
+        return false;
+    }
+    const repack = { labelType, mapBaseUrl };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/task/summary.js

@@ -0,0 +1,13 @@
+export function isSummary(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { preview, overview } = raw;
+    if (!(typeof preview === 'string' && typeof overview === 'string')) {
+        return false;
+    }
+    const repack = { preview, overview };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/{core/task → task}/task.d.ts

@@ -1,12 +1,12 @@
-import { DisplayInformation } from './display_information';
-import { TrainingInformation } from './training_information';
-import { Digest } from './digest';
-export declare type TaskID = string;
-export declare function isTaskID(obj: unknown): obj is TaskID;
-export declare function isTask(raw: unknown): raw is Task;
+import { type DisplayInformation } from './display_information.js';
+import { type TrainingInformation } from './training_information.js';
+import { type Digest } from './digest.js';
+export type TaskID = string;
 export interface Task {
-    taskID: TaskID;
+    id: TaskID;
     digest?: Digest;
     displayInformation: DisplayInformation;
     trainingInformation: TrainingInformation;
 }
+export declare function isTaskID(obj: unknown): obj is TaskID;
+export declare function isTask(raw: unknown): raw is Task;

package/dist/task/task.js

@@ -0,0 +1,22 @@
+import { isDisplayInformation } from './display_information.js';
+import { isTrainingInformation } from './training_information.js';
+import { isDigest } from './digest.js';
+export function isTaskID(obj) {
+    return typeof obj === 'string';
+}
+export function isTask(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { id, digest, displayInformation, trainingInformation } = raw;
+    if (!isTaskID(id) ||
+        (digest !== undefined && !isDigest(digest)) ||
+        !isDisplayInformation(displayInformation) ||
+        !isTrainingInformation(trainingInformation)) {
+        return false;
+    }
+    const repack = { id, digest, displayInformation, trainingInformation };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/task/task_handler.d.ts

@@ -0,0 +1,5 @@
+import { Map } from 'immutable';
+import type { Model } from '../index.js';
+import type { Task, TaskID } from './task.js';
+export declare function pushTask(url: URL, task: Task, model: Model): Promise<void>;
+export declare function fetchTasks(url: URL): Promise<Map<TaskID, Task>>;

package/dist/task/task_handler.js

@@ -0,0 +1,20 @@
+import axios from 'axios';
+import { Map } from 'immutable';
+import { serialization } from '../index.js';
+import { isTask } from './task.js';
+const TASK_ENDPOINT = 'tasks';
+export async function pushTask(url, task, model) {
+    await axios.post(url.href + TASK_ENDPOINT, {
+        task,
+        model: await serialization.model.encode(model),
+        weights: await serialization.weights.encode(model.weights)
+    });
+}
+export async function fetchTasks(url) {
+    const response = await axios.get(new URL(TASK_ENDPOINT, url).href);
+    const tasks = response.data;
+    if (!(Array.isArray(tasks) && tasks.every(isTask))) {
+        throw new Error('invalid tasks response');
+    }
+    return Map(tasks.map((t) => [t.id, t]));
+}

package/dist/task/task_provider.d.ts

@@ -0,0 +1,5 @@
+import type { Model, Task } from '../index.js';
+export interface TaskProvider {
+    getTask: () => Task;
+    getModel: () => Promise<Model>;
+}

package/dist/task/task_provider.js

@@ -0,0 +1 @@
+export {};

package/dist/{core/task → task}/training_information.d.ts

@@ -1,6 +1,6 @@
-import { Preprocessing } from '../dataset/data/preprocessing';
-import { ModelCompileData } from './model_compile_data';
-export declare function isTrainingInformation(raw: unknown): raw is TrainingInformation;
+import type { AggregatorChoice } from '../aggregator/get.js';
+import type { Preprocessing } from '../dataset/data/preprocessing/index.js';
+import { PreTrainedTokenizer } from '@xenova/transformers';
 export interface TrainingInformation {
     modelID: string;
     epochs: number;
@@ -8,21 +8,20 @@ export interface TrainingInformation {
     validationSplit: number;
     batchSize: number;
     preprocessingFunctions?: Preprocessing[];
-    modelCompileData: ModelCompileData;
-    dataType: string;
+    dataType: 'image' | 'tabular' | 'text';
     inputColumns?: string[];
     outputColumns?: string[];
     IMAGE_H?: number;
     IMAGE_W?: number;
-    modelURL?: string;
     LABEL_LIST?: string[];
-    learningRate?: number;
-    scheme: string;
+    scheme: 'decentralized' | 'federated' | 'local';
     noiseScale?: number;
     clippingRadius?: number;
     decentralizedSecure?: boolean;
-    byzantineRobustAggregator?: boolean;
-    tauPercentile?: number;
     maxShareValue?: number;
     minimumReadyPeers?: number;
+    aggregator?: AggregatorChoice;
+    tokenizer?: string | PreTrainedTokenizer;
+    maxSequenceLength?: number;
 }
+export declare function isTrainingInformation(raw: unknown): raw is TrainingInformation;

package/dist/task/training_information.js

@@ -0,0 +1,88 @@
+import { PreTrainedTokenizer } from '@xenova/transformers';
+function isStringArray(raw) {
+    if (!Array.isArray(raw)) {
+        return false;
+    }
+    const arr = raw; // isArray is unsafely guarding with any[]
+    return arr.every((e) => typeof e === 'string');
+}
+export function isTrainingInformation(raw) {
+    if (typeof raw !== 'object' || raw === null) {
+        return false;
+    }
+    const { IMAGE_H, IMAGE_W, LABEL_LIST, aggregator, batchSize, clippingRadius, dataType, decentralizedSecure, epochs, inputColumns, maxShareValue, minimumReadyPeers, modelID, noiseScale, outputColumns, preprocessingFunctions, roundDuration, scheme, validationSplit, tokenizer, maxSequenceLength, } = raw;
+    if (typeof dataType !== 'string' ||
+        typeof modelID !== 'string' ||
+        typeof epochs !== 'number' ||
+        typeof batchSize !== 'number' ||
+        typeof roundDuration !== 'number' ||
+        typeof validationSplit !== 'number' ||
+        (tokenizer !== undefined && typeof tokenizer !== 'string' && !(tokenizer instanceof PreTrainedTokenizer)) ||
+        (maxSequenceLength !== undefined && typeof maxSequenceLength !== 'number') ||
+        (aggregator !== undefined && typeof aggregator !== 'number') ||
+        (clippingRadius !== undefined && typeof clippingRadius !== 'number') ||
+        (decentralizedSecure !== undefined && typeof decentralizedSecure !== 'boolean') ||
+        (maxShareValue !== undefined && typeof maxShareValue !== 'number') ||
+        (minimumReadyPeers !== undefined && typeof minimumReadyPeers !== 'number') ||
+        (noiseScale !== undefined && typeof noiseScale !== 'number') ||
+        (IMAGE_H !== undefined && typeof IMAGE_H !== 'number') ||
+        (IMAGE_W !== undefined && typeof IMAGE_W !== 'number') ||
+        (LABEL_LIST !== undefined && !isStringArray(LABEL_LIST)) ||
+        (inputColumns !== undefined && !isStringArray(inputColumns)) ||
+        (outputColumns !== undefined && !isStringArray(outputColumns)) ||
+        (preprocessingFunctions !== undefined && !Array.isArray(preprocessingFunctions))) {
+        return false;
+    }
+    switch (dataType) {
+        case 'image': break;
+        case 'tabular': break;
+        case 'text': break;
+        default: return false;
+    }
+    // interdepences on data type
+    if (dataType === 'image') {
+        if (typeof IMAGE_H !== 'number' || typeof IMAGE_W !== 'number') {
+            return false;
+        }
+    }
+    else if (dataType in ['text', 'tabular']) {
+        if (!(Array.isArray(inputColumns) && inputColumns.every((e) => typeof e === 'string'))) {
+            return false;
+        }
+        if (!(Array.isArray(outputColumns) && outputColumns.every((e) => typeof e === 'string'))) {
+            return false;
+        }
+    }
+    switch (scheme) {
+        case 'decentralized': break;
+        case 'federated': break;
+        case 'local': break;
+        default: return false;
+    }
+    const repack = {
+        IMAGE_W,
+        IMAGE_H,
+        LABEL_LIST,
+        aggregator,
+        batchSize,
+        clippingRadius,
+        dataType,
+        decentralizedSecure,
+        epochs,
+        inputColumns,
+        maxShareValue,
+        minimumReadyPeers,
+        modelID,
+        noiseScale,
+        outputColumns,
+        preprocessingFunctions,
+        roundDuration,
+        scheme,
+        validationSplit,
+        tokenizer,
+        maxSequenceLength
+    };
+    const _correct = repack;
+    const _total = repack;
+    return true;
+}

package/dist/training/disco.d.ts

@@ -0,0 +1,40 @@
+import type { data, 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';
+export interface DiscoOptions {
+    client?: clients.Client;
+    aggregator?: Aggregator;
+    url?: string | URL;
+    scheme?: TrainingInformation['scheme'];
+    logger?: Logger;
+    memory?: Memory;
+}
+/**
+ * Top-level class handling distributed training from a client's perspective. It is meant to be
+ * a convenient object providing a reduced yet complete API that wraps model training,
+ * communication with nodes, logs and model memory.
+ */
+export declare class Disco {
+    readonly task: Task;
+    readonly logger: Logger;
+    readonly memory: Memory;
+    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 & {
+        participants: number;
+    }>;
+    /**
+     * Stops the ongoing training instance without disconnecting the client.
+     */
+    pause(): Promise<void>;
+    /**
+     * Completely stops the ongoing training instance.
+     */
+    close(): Promise<void>;
+}