@epfml/discojs 2.1.1 → 2.1.2-p20240506085037.0

package/dist/core/training/disco.js

@@ -1,130 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.Disco = void 0;
-var tslib_1 = require("tslib");
-var __1 = require("..");
-var trainer_builder_1 = require("./trainer/trainer_builder");
-// Handles the training loop, server communication & provides the user with feedback.
-var Disco = /** @class */ (function () {
-    // client need to be connected
-    function Disco(task, options) {
-        if (options.scheme === undefined) {
-            options.scheme = __1.TrainingSchemes[task.trainingInformation.scheme];
-        }
-        if (options.client === undefined) {
-            if (options.url === undefined) {
-                throw new Error('could not determine client from given parameters');
-            }
-            if (typeof options.url === 'string') {
-                options.url = new URL(options.url);
-            }
-            switch (options.scheme) {
-                case __1.TrainingSchemes.FEDERATED:
-                    options.client = new __1.client.federated.Client(options.url, task);
-                    break;
-                case __1.TrainingSchemes.DECENTRALIZED:
-                    options.client = new __1.client.federated.Client(options.url, task);
-                    break;
-                default:
-                    options.client = new __1.client.Local(options.url, task);
-                    break;
-            }
-        }
-        if (options.informant === undefined) {
-            switch (options.scheme) {
-                case __1.TrainingSchemes.FEDERATED:
-                    options.informant = new __1.informant.FederatedInformant(task);
-                    break;
-                case __1.TrainingSchemes.DECENTRALIZED:
-                    options.informant = new __1.informant.DecentralizedInformant(task);
-                    break;
-                default:
-                    options.informant = new __1.informant.LocalInformant(task);
-                    break;
-            }
-        }
-        if (options.logger === undefined) {
-            options.logger = new __1.ConsoleLogger();
-        }
-        if (options.memory === undefined) {
-            options.memory = new __1.EmptyMemory();
-        }
-        if (options.client.task !== task) {
-            throw new Error('client not setup for given task');
-        }
-        if (options.informant.task.taskID !== task.taskID) {
-            throw new Error('informant not setup for given task');
-        }
-        this.task = task;
-        this.client = options.client;
-        this.memory = options.memory;
-        this.logger = options.logger;
-        var trainerBuilder = new trainer_builder_1.TrainerBuilder(this.memory, this.task, options.informant);
-        this.trainer = trainerBuilder.build(this.client, options.scheme !== __1.TrainingSchemes.LOCAL);
-    }
-    Disco.prototype.fit = function (dataTuple) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var trainDataset, valDataset;
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0:
-                        this.logger.success('Thank you for your contribution. Data preprocessing has started');
-                        trainDataset = dataTuple.train.batch().preprocess();
-                        valDataset = dataTuple.validation !== undefined
-                            ? dataTuple.validation.batch().preprocess()
-                            : trainDataset;
-                        return [4 /*yield*/, this.client.connect()];
-                    case 1:
-                        _a.sent();
-                        return [4 /*yield*/, this.trainer];
-                    case 2: return [4 /*yield*/, (_a.sent()).trainModel(trainDataset.dataset, valDataset.dataset)];
-                    case 3:
-                        _a.sent();
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    // Stops the training function. Does not disconnect the client.
-    Disco.prototype.pause = function () {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0: return [4 /*yield*/, this.trainer];
-                    case 1: return [4 /*yield*/, (_a.sent()).stopTraining()];
-                    case 2:
-                        _a.sent();
-                        this.logger.success('Training was successfully interrupted.');
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    Disco.prototype.close = function () {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0: return [4 /*yield*/, this.pause()];
-                    case 1:
-                        _a.sent();
-                        return [4 /*yield*/, this.client.disconnect()];
-                    case 2:
-                        _a.sent();
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    Disco.prototype.logs = function () {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0: return [4 /*yield*/, this.trainer];
-                    case 1: return [2 /*return*/, (_a.sent()).getTrainerLog()];
-                }
-            });
-        });
-    };
-    return Disco;
-}());
-exports.Disco = Disco;

package/dist/core/training/index.d.ts

@@ -1,2 +0,0 @@
-export { Disco } from './disco';
-export { TrainingSchemes } from './training_schemes';

package/dist/core/training/index.js

@@ -1,7 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.TrainingSchemes = exports.Disco = void 0;
-var disco_1 = require("./disco");
-Object.defineProperty(exports, "Disco", { enumerable: true, get: function () { return disco_1.Disco; } });
-var training_schemes_1 = require("./training_schemes");
-Object.defineProperty(exports, "TrainingSchemes", { enumerable: true, get: function () { return training_schemes_1.TrainingSchemes; } });

package/dist/core/training/trainer/distributed_trainer.d.ts

@@ -1,20 +0,0 @@
-import { tf, Client, Memory, Task, TrainingInformant } from '../..';
-import { Trainer } from './trainer';
-/**
- * Class whose role is to train a model in a distributed way with a given dataset.
- */
-export declare class DistributedTrainer extends Trainer {
-    private readonly previousRoundModel;
-    private readonly client;
-    /** DistributedTrainer constructor, accepts same arguments as Trainer and in additional also a client who takes care of communicating weights.
-     */
-    constructor(task: Task, trainingInformant: TrainingInformant, memory: Memory, model: tf.LayersModel, previousRoundModel: tf.LayersModel, client: Client);
-    /**
-     * Callback called every time a round is over
-     */
-    onRoundEnd(accuracy: number): Promise<void>;
-    /**
-     * Callback called once training is over
-     */
-    onTrainEnd(): Promise<void>;
-}

package/dist/core/training/trainer/distributed_trainer.js

@@ -1,65 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.DistributedTrainer = void 0;
-var tslib_1 = require("tslib");
-var __1 = require("../..");
-var trainer_1 = require("./trainer");
-/**
- * Class whose role is to train a model in a distributed way with a given dataset.
- */
-var DistributedTrainer = /** @class */ (function (_super) {
-    (0, tslib_1.__extends)(DistributedTrainer, _super);
-    /** DistributedTrainer constructor, accepts same arguments as Trainer and in additional also a client who takes care of communicating weights.
-     */
-    function DistributedTrainer(task, trainingInformant, memory, model, previousRoundModel, client) {
-        var _this = _super.call(this, task, trainingInformant, memory, model) || this;
-        _this.previousRoundModel = previousRoundModel;
-        _this.client = client;
-        return _this;
-    }
-    /**
-     * Callback called every time a round is over
-     */
-    DistributedTrainer.prototype.onRoundEnd = function (accuracy) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var currentRoundWeights, previousRoundWeights, aggregatedWeights;
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0:
-                        currentRoundWeights = __1.WeightsContainer.from(this.model);
-                        previousRoundWeights = __1.WeightsContainer.from(this.previousRoundModel);
-                        return [4 /*yield*/, this.client.onRoundEndCommunication(currentRoundWeights, previousRoundWeights, this.roundTracker.round, this.trainingInformant)];
-                    case 1:
-                        aggregatedWeights = _a.sent();
-                        this.previousRoundModel.setWeights(currentRoundWeights.weights);
-                        this.model.setWeights(aggregatedWeights.weights);
-                        return [4 /*yield*/, this.memory.updateWorkingModel({ taskID: this.task.taskID, name: this.trainingInformation.modelID }, this.model)];
-                    case 2:
-                        _a.sent();
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    // if it is undefined, will training continue? we hope yes
-    /**
-     * Callback called once training is over
-     */
-    DistributedTrainer.prototype.onTrainEnd = function () {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0: return [4 /*yield*/, this.client.onTrainEndCommunication(__1.WeightsContainer.from(this.model), this.trainingInformant)];
-                    case 1:
-                        _a.sent();
-                        return [4 /*yield*/, _super.prototype.onTrainEnd.call(this)];
-                    case 2:
-                        _a.sent();
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    return DistributedTrainer;
-}(trainer_1.Trainer));
-exports.DistributedTrainer = DistributedTrainer;

package/dist/core/training/trainer/local_trainer.d.ts

@@ -1,11 +0,0 @@
-import { tf } from '../..';
-import { Trainer } from './trainer';
-/** Class whose role is to locally (alone) train a model on a given dataset, without any collaborators.
- */
-export declare class LocalTrainer extends Trainer {
-    /**
-     * Callback called every time a round is over. For local training, a round is typically an epoch
-     */
-    onRoundEnd(accuracy: number): Promise<void>;
-    protected onEpochEnd(epoch: number, logs?: tf.Logs): void;
-}

package/dist/core/training/trainer/local_trainer.js

@@ -1,34 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.LocalTrainer = void 0;
-var tslib_1 = require("tslib");
-var trainer_1 = require("./trainer");
-/** Class whose role is to locally (alone) train a model on a given dataset, without any collaborators.
- */
-var LocalTrainer = /** @class */ (function (_super) {
-    (0, tslib_1.__extends)(LocalTrainer, _super);
-    function LocalTrainer() {
-        return _super !== null && _super.apply(this, arguments) || this;
-    }
-    /**
-     * Callback called every time a round is over. For local training, a round is typically an epoch
-     */
-    LocalTrainer.prototype.onRoundEnd = function (accuracy) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0: return [4 /*yield*/, this.memory.updateWorkingModel({ taskID: this.task.taskID, name: this.trainingInformation.modelID }, this.model)];
-                    case 1:
-                        _a.sent();
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    LocalTrainer.prototype.onEpochEnd = function (epoch, logs) {
-        _super.prototype.onEpochEnd.call(this, epoch, logs);
-        this.trainingInformant.update({ currentRound: epoch });
-    };
-    return LocalTrainer;
-}(trainer_1.Trainer));
-exports.LocalTrainer = LocalTrainer;

package/dist/core/training/trainer/round_tracker.d.ts

@@ -1,30 +0,0 @@
-/**
- * Class that keeps track of the current batch in order for the trainer to query when round has ended.
- *
- * @remark
- * In distributed training, the client trains locally for a certain amount of epochs before sharing his weights to the server/neighbor, this
- * is what we call a round.
- *
- * The role of the RoundTracker is to keep track of when a roundHasEnded using the current batch number. The batch in the RoundTracker is cumulative whereas
- * in the onBatchEnd it is not (it resets to 0 after each epoch).
- *
- * The roundDuration is the length of a round (in batches).
- */
-export declare class RoundTracker {
-    round: number;
-    batch: number;
-    roundDuration: number;
-    constructor(roundDuration: number);
-    /**
-     * Update the batch number, to be called inside onBatchEnd. (We do not use batch output of onBatchEnd since it is
-     * not cumulative).
-     */
-    updateBatch(): void;
-    /**
-     * Returns true if a local round has ended, false otherwise.
-     *
-     * @remark
-     * Returns true if (batch) mod (batches per round) == 0, false otherwise
-     */
-    roundHasEnded(): boolean;
-}

package/dist/core/training/trainer/round_tracker.js

@@ -1,47 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.RoundTracker = void 0;
-/**
- * Class that keeps track of the current batch in order for the trainer to query when round has ended.
- *
- * @remark
- * In distributed training, the client trains locally for a certain amount of epochs before sharing his weights to the server/neighbor, this
- * is what we call a round.
- *
- * The role of the RoundTracker is to keep track of when a roundHasEnded using the current batch number. The batch in the RoundTracker is cumulative whereas
- * in the onBatchEnd it is not (it resets to 0 after each epoch).
- *
- * The roundDuration is the length of a round (in batches).
- */
-var RoundTracker = /** @class */ (function () {
-    function RoundTracker(roundDuration) {
-        this.round = 0;
-        this.batch = 0;
-        this.roundDuration = roundDuration;
-    }
-    /**
-     * Update the batch number, to be called inside onBatchEnd. (We do not use batch output of onBatchEnd since it is
-     * not cumulative).
-     */
-    RoundTracker.prototype.updateBatch = function () {
-        this.batch += 1;
-    };
-    /**
-     * Returns true if a local round has ended, false otherwise.
-     *
-     * @remark
-     * Returns true if (batch) mod (batches per round) == 0, false otherwise
-     */
-    RoundTracker.prototype.roundHasEnded = function () {
-        if (this.batch === 0) {
-            return false;
-        }
-        var roundHasEnded = this.batch % this.roundDuration === 0;
-        if (roundHasEnded) {
-            this.round += 1;
-        }
-        return roundHasEnded;
-    };
-    return RoundTracker;
-}());
-exports.RoundTracker = RoundTracker;

package/dist/core/training/trainer/trainer.d.ts

@@ -1,65 +0,0 @@
-import { tf, Memory, Task, TrainingInformant, TrainingInformation } from '../..';
-import { RoundTracker } from './round_tracker';
-import { TrainerLog } from '../../logging/trainer_logger';
-/** Abstract class whose role is to train a model with a given dataset. This can be either done
- * locally (alone) or in a distributed way with collaborators. The Trainer works as follows:
- *
- * 1. Call trainModel(dataset) to start training
- * 2. Once a batch ends, onBatchEnd is triggered, which will then call onRoundEnd once the round has ended.
- *
- * The onRoundEnd needs to be implemented to specify what actions to do when the round has ended, such as a communication step with collaborators. To know when
- * a round has ended we use the roundTracker object.
- */
-export declare abstract class Trainer {
-    readonly task: Task;
-    readonly trainingInformant: TrainingInformant;
-    readonly memory: Memory;
-    readonly model: tf.LayersModel;
-    readonly trainingInformation: TrainingInformation;
-    readonly roundTracker: RoundTracker;
-    private stopTrainingRequested;
-    private readonly trainerLogger;
-    /**
-     * Constructs the training manager.
-     * @param task the trained task
-     * @param trainingInformant the training informant
-     */
-    constructor(task: Task, trainingInformant: TrainingInformant, memory: Memory, model: tf.LayersModel);
-    /**
-     * Every time a round ends this function will be called
-     */
-    protected abstract onRoundEnd(accuracy: number): Promise<void>;
-    /** onBatchEnd callback, when a round ends, we call onRoundEnd (to be implemented for local and distributed instances)
-     */
-    protected onBatchEnd(_: number, logs?: tf.Logs): Promise<void>;
-    /**
-     * We update the training graph, this needs to be done on epoch end as there is no validation accuracy onBatchEnd.
-     */
-    protected onEpochEnd(epoch: number, logs?: tf.Logs): void;
-    /**
-     * When the training ends this function will be call
-     */
-    protected onTrainEnd(logs?: tf.Logs): Promise<void>;
-    /**
-     * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
-     */
-    stopTraining(): Promise<void>;
-    /**
-     * Start training the model with the given dataset
-     * @param dataset
-     */
-    trainModel(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): Promise<void>;
-    /**
-     * Format accuracy
-     */
-    protected roundDecimals(accuracy: number, decimalsToRound?: number): number;
-    /**
-     * reset stop training state
-     */
-    protected resetStopTrainerState(): void;
-    /**
-     * If stop training is requested, do so
-     */
-    protected stopTrainModelIfRequested(): void;
-    getTrainerLog(): TrainerLog;
-}

package/dist/core/training/trainer/trainer.js

@@ -1,160 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.Trainer = void 0;
-var tslib_1 = require("tslib");
-var round_tracker_1 = require("./round_tracker");
-var trainer_logger_1 = require("../../logging/trainer_logger");
-/** Abstract class whose role is to train a model with a given dataset. This can be either done
- * locally (alone) or in a distributed way with collaborators. The Trainer works as follows:
- *
- * 1. Call trainModel(dataset) to start training
- * 2. Once a batch ends, onBatchEnd is triggered, which will then call onRoundEnd once the round has ended.
- *
- * The onRoundEnd needs to be implemented to specify what actions to do when the round has ended, such as a communication step with collaborators. To know when
- * a round has ended we use the roundTracker object.
- */
-var Trainer = /** @class */ (function () {
-    /**
-     * Constructs the training manager.
-     * @param task the trained task
-     * @param trainingInformant the training informant
-     */
-    function Trainer(task, trainingInformant, memory, model) {
-        this.task = task;
-        this.trainingInformant = trainingInformant;
-        this.memory = memory;
-        this.model = model;
-        this.stopTrainingRequested = false;
-        this.trainerLogger = new trainer_logger_1.TrainerLogger();
-        var trainingInformation = task.trainingInformation;
-        if (trainingInformation === undefined) {
-            throw new Error('round duration is undefined');
-        }
-        this.trainingInformation = trainingInformation;
-        this.roundTracker = new round_tracker_1.RoundTracker(trainingInformation.roundDuration);
-    }
-    /** onBatchEnd callback, when a round ends, we call onRoundEnd (to be implemented for local and distributed instances)
-     */
-    Trainer.prototype.onBatchEnd = function (_, logs) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0:
-                        if (logs === undefined) {
-                            return [2 /*return*/];
-                        }
-                        this.roundTracker.updateBatch();
-                        this.stopTrainModelIfRequested();
-                        if (!this.roundTracker.roundHasEnded()) return [3 /*break*/, 2];
-                        return [4 /*yield*/, this.onRoundEnd(logs.acc)];
-                    case 1:
-                        _a.sent();
-                        _a.label = 2;
-                    case 2: return [2 /*return*/];
-                }
-            });
-        });
-    };
-    /**
-     * We update the training graph, this needs to be done on epoch end as there is no validation accuracy onBatchEnd.
-     */
-    Trainer.prototype.onEpochEnd = function (epoch, logs) {
-        this.trainerLogger.onEpochEnd(epoch, logs);
-        if (logs !== undefined && !isNaN(logs.acc) && !isNaN(logs.val_acc)) {
-            this.trainingInformant.updateTrainingGraph(this.roundDecimals(logs.acc));
-            this.trainingInformant.updateValidationGraph(this.roundDecimals(logs.val_acc));
-        }
-        else {
-            this.trainerLogger.error('onEpochEnd: NaN value');
-        }
-    };
-    /**
-     * When the training ends this function will be call
-     */
-    Trainer.prototype.onTrainEnd = function (logs) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                this.trainingInformant.addMessage('Training finished.');
-                return [2 /*return*/];
-            });
-        });
-    };
-    /**
-     * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
-     */
-    Trainer.prototype.stopTraining = function () {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                this.stopTrainingRequested = true;
-                return [2 /*return*/];
-            });
-        });
-    };
-    /**
-     * Start training the model with the given dataset
-     * @param dataset
-     */
-    Trainer.prototype.trainModel = function (dataset, valDataset) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var _this = this;
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0:
-                        this.resetStopTrainerState();
-                        // Assign callbacks and start training
-                        return [4 /*yield*/, this.model.fitDataset(dataset, {
-                                epochs: this.trainingInformation.epochs,
-                                validationData: valDataset,
-                                callbacks: {
-                                    onEpochEnd: function (epoch, logs) { return _this.onEpochEnd(epoch, logs); },
-                                    onBatchEnd: function (epoch, logs) { return (0, tslib_1.__awaiter)(_this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {
-                                        switch (_a.label) {
-                                            case 0: return [4 /*yield*/, this.onBatchEnd(epoch, logs)];
-                                            case 1: return [2 /*return*/, _a.sent()];
-                                        }
-                                    }); }); },
-                                    onTrainEnd: function (logs) { return (0, tslib_1.__awaiter)(_this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {
-                                        switch (_a.label) {
-                                            case 0: return [4 /*yield*/, this.onTrainEnd(logs)];
-                                            case 1: return [2 /*return*/, _a.sent()];
-                                        }
-                                    }); }); }
-                                }
-                            })];
-                    case 1:
-                        // Assign callbacks and start training
-                        _a.sent();
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    /**
-     * Format accuracy
-     */
-    Trainer.prototype.roundDecimals = function (accuracy, decimalsToRound) {
-        if (decimalsToRound === void 0) { decimalsToRound = 2; }
-        return +(accuracy * 100).toFixed(decimalsToRound);
-    };
-    /**
-     * reset stop training state
-     */
-    Trainer.prototype.resetStopTrainerState = function () {
-        this.model.stopTraining = false;
-        this.stopTrainingRequested = false;
-    };
-    /**
-     * If stop training is requested, do so
-     */
-    Trainer.prototype.stopTrainModelIfRequested = function () {
-        if (this.stopTrainingRequested) {
-            this.model.stopTraining = true;
-            this.stopTrainingRequested = false;
-        }
-    };
-    Trainer.prototype.getTrainerLog = function () {
-        return this.trainerLogger.log;
-    };
-    return Trainer;
-}());
-exports.Trainer = Trainer;

package/dist/core/training/trainer/trainer_builder.js

@@ -1,95 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.TrainerBuilder = void 0;
-var tslib_1 = require("tslib");
-var __1 = require("../..");
-var distributed_trainer_1 = require("./distributed_trainer");
-var local_trainer_1 = require("./local_trainer");
-/**
- * A class that helps build the Trainer and auxiliary classes.
- */
-var TrainerBuilder = /** @class */ (function () {
-    function TrainerBuilder(memory, task, trainingInformant) {
-        this.memory = memory;
-        this.task = task;
-        this.trainingInformant = trainingInformant;
-    }
-    /**
-     * Builds a trainer object.
-     *
-     * @param client client to share weights with (either distributed or federated)
-     * @param distributed whether to build a distributed or local trainer
-     * @returns
-     */
-    TrainerBuilder.prototype.build = function (client, distributed) {
-        if (distributed === void 0) { distributed = false; }
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var model;
-            return (0, tslib_1.__generator)(this, function (_a) {
-                switch (_a.label) {
-                    case 0: return [4 /*yield*/, this.getModel(client)];
-                    case 1:
-                        model = _a.sent();
-                        if (distributed) {
-                            return [2 /*return*/, new distributed_trainer_1.DistributedTrainer(this.task, this.trainingInformant, this.memory, model, model, client)];
-                        }
-                        else {
-                            return [2 /*return*/, new local_trainer_1.LocalTrainer(this.task, this.trainingInformant, this.memory, model)];
-                        }
-                        return [2 /*return*/];
-                }
-            });
-        });
-    };
-    /**
-     * If a model exists in memory, laod it, otherwise load model from server
-     * @returns
-     */
-    TrainerBuilder.prototype.getModel = function (client) {
-        var _a;
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var modelID, info, model;
-            return (0, tslib_1.__generator)(this, function (_b) {
-                switch (_b.label) {
-                    case 0:
-                        modelID = (_a = this.task.trainingInformation) === null || _a === void 0 ? void 0 : _a.modelID;
-                        if (modelID === undefined) {
-                            throw new TypeError('model ID is undefined');
-                        }
-                        info = { type: __1.ModelType.WORKING, taskID: this.task.taskID, name: modelID };
-                        return [4 /*yield*/, this.memory.contains(info)];
-                    case 1: return [4 /*yield*/, ((_b.sent()) ? this.memory.getModel(info) : client.getLatestModel())];
-                    case 2:
-                        model = _b.sent();
-                        return [4 /*yield*/, this.updateModelInformation(model)];
-                    case 3: return [2 /*return*/, _b.sent()];
-                }
-            });
-        });
-    };
-    TrainerBuilder.prototype.updateModelInformation = function (model) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var info;
-            return (0, tslib_1.__generator)(this, function (_a) {
-                // Continue local training from previous epoch checkpoint
-                if (model.getUserDefinedMetadata() === undefined) {
-                    model.setUserDefinedMetadata({ epoch: 0 });
-                }
-                info = this.task.trainingInformation;
-                if (info === undefined) {
-                    throw new TypeError('training information is undefined');
-                }
-                model.compile(info.modelCompileData);
-                if (info.learningRate !== undefined) {
-                    // TODO: Not the right way to change learningRate and hence we cast to any
-                    // the right way is to construct the optimiser and pass learningRate via
-                    // argument.
-                    model.optimizer.learningRate = info.learningRate;
-                }
-                return [2 /*return*/, model];
-            });
-        });
-    };
-    return TrainerBuilder;
-}());
-exports.TrainerBuilder = TrainerBuilder;