@epfml/discojs 3.0.1-p20240821133014.0 → 3.0.1-p20240826092658.0

package/dist/dataset/data/data.d.ts

@@ -1,23 +1,22 @@
 import * as tf from '@tensorflow/tfjs';
 import type { List } from 'immutable';
 import type { Task } from '../../index.js';
-import type { Dataset } from '../index.js';
 import type { PreprocessingFunction } from './preprocessing/base.js';
 /**
  * Abstract class representing an immutable Disco dataset, including a TF.js dataset,
  * Disco task and set of preprocessing functions.
  */
 export declare abstract class Data {
-    readonly dataset: Dataset;
+    readonly dataset: tf.data.Dataset<tf.TensorContainer>;
     readonly task: Task;
     readonly size?: number | undefined;
     abstract readonly availablePreprocessing: List<PreprocessingFunction>;
-    protected constructor(dataset: Dataset, task: Task, size?: number | undefined);
-    static init(_dataset: Dataset, _task: Task, _size?: number): Promise<Data>;
+    protected constructor(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size?: number | undefined);
+    static init(_dataset: tf.data.Dataset<tf.TensorContainer>, _task: Task, _size?: number): Promise<Data>;
     /**
      * Callable abstract method instead of constructor.
      */
-    protected abstract create(dataset: Dataset, task: Task, size?: number): Data;
+    protected abstract create(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size?: number): Data;
     /**
      * Creates a new Disco data object containing the batched TF.js dataset, according to the
      * task's parameters.
@@ -27,7 +26,7 @@ export declare abstract class Data {
     /**
      * The TF.js dataset batched according to the task's parameters.
      */
-    get batchedDataset(): Dataset;
+    get batchedDataset(): tf.data.Dataset<tf.TensorContainer>;
     /**
      * Creates a new Disco data object containing the preprocessed TF.js dataset,
      * according to the defined set of preprocessing functions and the task's parameters.
@@ -44,5 +43,5 @@ export declare abstract class Data {
      * The TF.js dataset preprocessing according to the set of preprocessing functions and the task's
      * parameters.
      */
-    get preprocessedDataset(): Dataset;
+    get preprocessedDataset(): tf.data.Dataset<tf.TensorContainer>;
 }

package/dist/dataset/data/data.js

@@ -57,17 +57,22 @@ export class Data {
         const applyPreprocessing = this.availablePreprocessing
             .filter((e) => e.type in taskPreprocessing)
             .map((e) => e.apply);
-        if (applyPreprocessing.size === 0) {
-            return x => Promise.resolve(x);
-        }
         const preprocessingChain = async (input) => {
             let currentContainer = await input; // Start with the initial tensor container
             for (const fn of applyPreprocessing) {
-                const newContainer = await fn(Promise.resolve(currentContainer), this.task);
-                if (currentContainer !== newContainer) {
-                    tf.dispose(currentContainer); // Dispose of the old container
+                const next = await fn(Promise.resolve(currentContainer), this.task);
+                // dirty but kinda working way to dispose of converted tensors
+                if (typeof currentContainer === "object" && typeof next === "object") {
+                    if ("xs" in currentContainer &&
+                        "xs" in next &&
+                        currentContainer.xs !== next.xs)
+                        tf.dispose(currentContainer.xs);
+                    if ("ys" in currentContainer &&
+                        "ys" in next &&
+                        currentContainer.ys !== next.ys)
+                        tf.dispose(currentContainer.ys);
                 }
-                currentContainer = newContainer;
+                currentContainer = next;
             }
             return currentContainer; // Return the final tensor container
         };

package/dist/dataset/data/helpers.d.ts

@@ -0,0 +1,10 @@
+/** Internal functions to help with Dataset to Data/DataSplit conversion
+ *
+ * @todo rm when fully using Dataset
+ */
+import type { Task, TypedDataset, TypedLabeledDataset } from "../../index.js";
+import { Data } from "./index.js";
+import { DataSplit } from "./data_split.js";
+export declare function datasetToData(task: Task, [t, dataset]: TypedDataset): Promise<Data>;
+export declare function labeledDatasetToData(task: Task, [t, dataset]: TypedLabeledDataset): Promise<Data>;
+export declare function labeledDatasetToDataSplit(task: Task, [t, dataset]: TypedLabeledDataset): Promise<DataSplit>;

package/dist/dataset/data/helpers.js

@@ -0,0 +1,97 @@
+/** Internal functions to help with Dataset to Data/DataSplit conversion
+ *
+ * @todo rm when fully using Dataset
+ */
+import { List } from "immutable";
+import * as tf from "@tensorflow/tfjs";
+import { processing } from "../../index.js";
+import { ImageData, TabularData, TextData } from "./index.js";
+function intoTFDataset(iter) {
+    // @ts-expect-error generator
+    return tf.data.generator(async function* () {
+        yield* iter;
+    });
+}
+function imageToTensor(image) {
+    return tf.tensor3d(image.data, [image.width, image.height, 3], "int32");
+}
+function tabularToNumbers(columns, row) {
+    return List(columns)
+        .map((column) => processing.extractColumn(row, column))
+        .map((v) => (v !== "" ? v : "0")) // TODO how to specify defaults?
+        .map(processing.convertToNumber)
+        .toArray();
+}
+export async function datasetToData(task, [t, dataset]) {
+    switch (t) {
+        case "image": {
+            const converted = dataset
+                .map(processing.removeAlpha)
+                .map((image) => processing.expandToMulticolor(image))
+                .map((image) => ({
+                xs: imageToTensor(image),
+            }));
+            return await ImageData.init(intoTFDataset(converted), task);
+        }
+        case "tabular": {
+            const inputColumns = task.trainingInformation.inputColumns;
+            if (inputColumns === undefined)
+                throw new Error("tabular task without input columns");
+            const converted = dataset.map((row) => ({
+                xs: tabularToNumbers(inputColumns, row),
+            }));
+            return await TabularData.init(intoTFDataset(converted), task);
+        }
+        case "text":
+            return await TextData.init(intoTFDataset(dataset), task);
+    }
+}
+export async function labeledDatasetToData(task, [t, dataset]) {
+    switch (t) {
+        case "image": {
+            const labels = List(task.trainingInformation.LABEL_LIST);
+            const converted = dataset
+                .map(([image, label]) => [
+                processing.expandToMulticolor(processing.removeAlpha(image)),
+                processing.indexInList(label, labels),
+            ])
+                .map(([image, label]) => ({
+                xs: imageToTensor(image),
+                ys: tf.oneHot(label, labels.size, 1, 0, "int32"),
+            }));
+            return await ImageData.init(intoTFDataset(converted), task);
+        }
+        case "tabular": {
+            const { inputColumns, outputColumns } = task.trainingInformation;
+            if (inputColumns === undefined || outputColumns === undefined)
+                throw new Error("tabular task without input and output columns");
+            const converted = dataset.map((row) => ({
+                xs: tabularToNumbers(inputColumns, row),
+                ys: tf.tensor1d(tabularToNumbers(outputColumns, row)),
+            }));
+            return await TabularData.init(intoTFDataset(converted), task);
+        }
+        case "text":
+            return await TextData.init(intoTFDataset(dataset), task);
+    }
+}
+export async function labeledDatasetToDataSplit(task, [t, dataset]) {
+    const split = task.trainingInformation.validationSplit;
+    let train;
+    let validation;
+    switch (t) {
+        case "image": {
+            [train, validation] = await Promise.all(dataset.split(split).map((d) => labeledDatasetToData(task, [t, d])));
+            break;
+        }
+        case "tabular": {
+            [train, validation] = await Promise.all(dataset.split(split).map((d) => labeledDatasetToData(task, [t, d])));
+            break;
+        }
+        case "text": {
+            [train, validation] = await Promise.all(dataset.split(split).map((d) => labeledDatasetToData(task, [t, d])));
+            break;
+        }
+    }
+    return { train, validation };
+}

package/dist/dataset/data/image_data.d.ts

@@ -1,11 +1,11 @@
+import * as tf from '@tensorflow/tfjs';
 import type { Task } from '../../index.js';
-import type { Dataset } from '../dataset.js';
 import { Data } from './data.js';
 /**
  * Disco data made of image samples (.jpg, .png, etc.).
  */
 export declare class ImageData extends Data {
     readonly availablePreprocessing: import("immutable").List<import("./preprocessing/base.js").PreprocessingFunction>;
-    static init(dataset: Dataset, task: Task, size?: number): Promise<Data>;
-    protected create(dataset: Dataset, task: Task, size: number): ImageData;
+    static init(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size?: number): Promise<Data>;
+    protected create(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size: number): ImageData;
 }

package/dist/dataset/data/image_data.js

@@ -1,3 +1,4 @@
+import * as tf from '@tensorflow/tfjs';
 import { Data } from './data.js';
 import { ImagePreprocessing, IMAGE_PREPROCESSING } from './preprocessing/index.js';
 /**
@@ -11,14 +12,17 @@ export class ImageData extends Data {
         // cause an error during training, because of the lazy aspect of the dataset; we only
         // verify the first sample.
         if (task.trainingInformation.preprocessingFunctions?.includes(ImagePreprocessing.Resize) !== true) {
-            const sample = (await dataset.take(1).toArray())[0];
+            const iteration = await dataset.iterator().then((iter) => iter.next());
+            if (iteration.done === true)
+                throw new Error("empty dataset");
+            const sample = iteration.value;
             // TODO: We suppose the presence of labels
             // TODO: Typing (discojs-node/src/dataset/data_loader/image_loader.spec.ts)
             if (typeof sample !== 'object' || sample === null || sample === undefined) {
                 throw new Error("Image is undefined or is not an object");
             }
             let shape;
-            if ('xs' in sample && 'ys' in sample) {
+            if ('xs' in sample) {
                 shape = sample.xs.shape;
             }
             else {
@@ -29,6 +33,7 @@ export class ImageData extends Data {
                 (shape[0] !== IMAGE_W || shape[1] !== IMAGE_H)) {
                 throw new Error(`Image doesn't have the dimensions specified in the task's training information. Expected ${IMAGE_H}x${IMAGE_W} but got ${shape[0]}x${shape[1]}.`);
             }
+            tf.dispose(sample);
         }
         return new ImageData(dataset, task, size);
     }

package/dist/dataset/data/index.d.ts

@@ -1,4 +1,3 @@
-export type { DataSplit } from './data_split.js';
 export { Data } from './data.js';
 export { ImageData } from './image_data.js';
 export { TabularData } from './tabular_data.js';

package/dist/dataset/data/preprocessing/text_preprocessing.js

@@ -9,6 +9,21 @@ export var TextPreprocessing;
     TextPreprocessing[TextPreprocessing["Tokenize"] = 0] = "Tokenize";
     TextPreprocessing[TextPreprocessing["LeftPadding"] = 1] = "LeftPadding";
 })(TextPreprocessing || (TextPreprocessing = {}));
+function isNumberArray(raw) {
+    if (!Array.isArray(raw))
+        return false;
+    const arr = raw; // isArray is unsafely guarding with any[]
+    return arr.every((e) => typeof e === "number");
+}
+function isTokenizedEntry(raw) {
+    if (typeof raw !== "object" || raw === null)
+        return false;
+    const { tokens } = raw;
+    if (!isNumberArray(tokens))
+        return false;
+    const _ = { tokens };
+    return true;
+}
 /**
  * LeftPadding pads all incoming inputs to be a fixed length, which should be specified
  * in `task.trainingInformation.maxSequenceLength`.
@@ -23,11 +38,11 @@ export var TextPreprocessing;
  */
 const leftPadding = {
     type: TextPreprocessing.LeftPadding,
-    apply: async (x, task) => {
-        if (x === undefined || !Array.isArray(x) || x.length == 0 || typeof (x[0] !== 'number')) {
-            new Error("The leftPadding preprocessing expects a non empty 1D array of number");
-        }
-        const { tokens } = await x;
+    apply: async (input, task) => {
+        const x = await input;
+        if (!isTokenizedEntry(x))
+            throw new Error("The leftPadding preprocessing expects a non empty 1D array of number");
+        const { tokens } = x;
         const tokenizer = await models.getTaskTokenizer(task);
         return tf.tidy(() => {
             // maxLength is the final length of xs
@@ -59,10 +74,9 @@ const leftPadding = {
 const tokenize = {
     type: TextPreprocessing.Tokenize,
     apply: async (x, task) => {
-        if (typeof x !== 'string') {
-            new Error("The tokenize preprocessing expects a string as input");
-        }
-        const xs = await x; // tf.TextLineDataset yields strings
+        const xs = await x;
+        if (typeof xs !== 'string')
+            throw new Error("The tokenize preprocessing expects a string as input");
         const tokenizer = await models.getTaskTokenizer(task);
         // Add plus one to include the next token label of the last token in the input sequence
         // The inputs are truncated down to exactly maxSequenceLength in leftPadding

package/dist/dataset/data/tabular_data.d.ts

@@ -1,11 +1,11 @@
+import * as tf from '@tensorflow/tfjs';
 import type { Task } from '../../index.js';
-import type { Dataset } from '../dataset.js';
 import { Data } from './data.js';
 /**
  * Disco data made of tabular (.csv, .tsv, etc.) files.
  */
 export declare class TabularData extends Data {
     readonly availablePreprocessing: import("immutable").List<import("./preprocessing/base.js").PreprocessingFunction>;
-    static init(dataset: Dataset, task: Task, size?: number): Promise<TabularData>;
-    protected create(dataset: Dataset, task: Task, size: number): TabularData;
+    static init(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size?: number): Promise<TabularData>;
+    protected create(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size: number): TabularData;
 }

package/dist/dataset/data/text_data.d.ts

@@ -1,11 +1,11 @@
+import * as tf from '@tensorflow/tfjs';
 import type { Task } from '../../index.js';
-import type { Dataset } from '../dataset.js';
 import { Data } from './data.js';
 /**
  * Disco data made of textual samples.
  */
 export declare class TextData extends Data {
     readonly availablePreprocessing: import("immutable").List<import("./preprocessing/base.js").PreprocessingFunction>;
-    static init(dataset: Dataset, task: Task, size?: number): Promise<TextData>;
-    protected create(dataset: Dataset, task: Task, size?: number): TextData;
+    static init(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size?: number): Promise<TextData>;
+    protected create(dataset: tf.data.Dataset<tf.TensorContainer>, task: Task, size?: number): TextData;
 }

package/dist/dataset/dataset.d.ts

@@ -1,5 +1,48 @@
-import type tf from '@tensorflow/tfjs';
-/**
- * Convenient type for the common dataset type used in TF.js.
- */
-export type Dataset = tf.data.Dataset<tf.TensorContainer>;
+import { List } from "immutable";
+type DatasetLike<T> = AsyncIterable<T> | Iterable<T> | (() => AsyncIterator<T, void>) | (() => Iterator<T, void>);
+/** Immutable series of data */
+export declare class Dataset<T> implements AsyncIterable<T> {
+    #private;
+    /** Wrap given data generator
+     *
+     * To avoid loading everything in memory, it is a function that upon calling
+     * should return a new AsyncGenerator with the same data as before.
+     */
+    constructor(content: DatasetLike<T>);
+    [Symbol.asyncIterator](): AsyncIterator<T>;
+    /** Apply function to each element
+     *
+     * @param mapper how to change each element
+     */
+    map<U>(mapper: (_: T) => U | Promise<U>): Dataset<U>;
+    /** Combine with another Dataset.
+     *
+     * @param other what to yield after us
+     */
+    chain(other: Dataset<T> | DatasetLike<T>): Dataset<T>;
+    /** Divide into two based on given ratio
+     *
+     * @param ratio between 0 (all on left) and 1 (all on right)
+     */
+    split(ratio: number): [Dataset<T>, Dataset<T>];
+    /** Slice into chunks
+     *
+     * Last slice is smaller if dataset isn't perfectly divisible
+     *
+     * @param size count of element per chunk
+     */
+    batch(size: number): Dataset<List<T>>;
+    /** Join side-by-side
+     *
+     * Stops as soon as one runs out
+     *
+     * @param other right side
+     **/
+    zip<U>(other: Dataset<U> | DatasetLike<U>): Dataset<[T, U]>;
+    /** Compute size
+     *
+     * This is a costly operation as we need to go through the whole Dataset.
+     */
+    size(): Promise<number>;
+}
+export {};

package/dist/dataset/dataset.js

@@ -1 +1,155 @@
-export {};
+import { List } from "immutable";
+/** Immutable series of data */
+export class Dataset {
+    #content;
+    /** Wrap given data generator
+     *
+     * To avoid loading everything in memory, it is a function that upon calling
+     * should return a new AsyncGenerator with the same data as before.
+     */
+    constructor(content) {
+        this.#content = async function* () {
+            let iter;
+            if (typeof content === "function")
+                iter = content();
+            else if (Symbol.asyncIterator in content)
+                iter = content[Symbol.asyncIterator]();
+            else
+                iter = content[Symbol.iterator]();
+            while (true) {
+                const result = await iter.next();
+                if (result.done === true)
+                    break;
+                yield result.value;
+            }
+        };
+    }
+    [Symbol.asyncIterator]() {
+        return this.#content();
+    }
+    /** Apply function to each element
+     *
+     * @param mapper how to change each element
+     */
+    map(mapper) {
+        const content = {
+            [Symbol.asyncIterator]: () => this.#content(),
+        };
+        return new Dataset(async function* () {
+            for await (const e of content)
+                yield await mapper(e);
+        });
+    }
+    /** Combine with another Dataset.
+     *
+     * @param other what to yield after us
+     */
+    chain(other) {
+        if (!(other instanceof Dataset))
+            other = new Dataset(other);
+        const self = {
+            [Symbol.asyncIterator]: () => this.#content(),
+        };
+        return new Dataset(async function* () {
+            yield* self;
+            yield* other;
+        });
+    }
+    /** Divide into two based on given ratio
+     *
+     * @param ratio between 0 (all on left) and 1 (all on right)
+     */
+    split(ratio) {
+        if (ratio < 0 || ratio > 1)
+            throw new Error("ratio out of range");
+        const content = {
+            [Symbol.asyncIterator]: () => this.#content(),
+        };
+        // to avoid using random sampling or knowing the size beforehand,
+        // we compute the actual ratio and make it converge towards the wanted one
+        return [
+            new Dataset(async function* () {
+                let yielded_by_other = 0;
+                let total_size = 0;
+                for await (const e of content) {
+                    total_size++;
+                    if (yielded_by_other / total_size >= ratio) {
+                        yield e;
+                    }
+                    else {
+                        yielded_by_other++;
+                    }
+                }
+            }),
+            new Dataset(async function* () {
+                let yielded = 0;
+                let total_size = 0;
+                for await (const e of content) {
+                    total_size++;
+                    if (yielded / total_size < ratio) {
+                        yielded++;
+                        yield e;
+                    }
+                }
+            }),
+        ];
+    }
+    /** Slice into chunks
+     *
+     * Last slice is smaller if dataset isn't perfectly divisible
+     *
+     * @param size count of element per chunk
+     */
+    batch(size) {
+        if (size <= 0 || !Number.isInteger(size))
+            throw new Error("invalid size");
+        const content = {
+            [Symbol.asyncIterator]: () => this.#content(),
+        };
+        return new Dataset(async function* () {
+            let batch = List();
+            for await (const e of content) {
+                batch = batch.push(e);
+                if (batch.size === size) {
+                    yield batch;
+                    batch = List();
+                }
+            }
+            if (!batch.isEmpty())
+                yield batch;
+        });
+    }
+    /** Join side-by-side
+     *
+     * Stops as soon as one runs out
+     *
+     * @param other right side
+     **/
+    zip(other) {
+        if (!(other instanceof Dataset))
+            other = new Dataset(other);
+        const content = {
+            [Symbol.asyncIterator]: () => this.#content(),
+        };
+        return new Dataset(async function* () {
+            const left = content[Symbol.asyncIterator]();
+            const right = other[Symbol.asyncIterator]();
+            while (true) {
+                const [l, r] = await Promise.all([left.next(), right.next()]);
+                if (l.done || r.done)
+                    return;
+                yield [l.value, r.value];
+            }
+        });
+    }
+    /** Compute size
+     *
+     * This is a costly operation as we need to go through the whole Dataset.
+     */
+    async size() {
+        let ret = 0;
+        for await (const _ of this)
+            ret++;
+        return ret;
+    }
+}

package/dist/dataset/image.d.ts

@@ -0,0 +1,14 @@
+/**
+ * Raw image with type level dimensions.
+ *
+ * @typeParam D depth of the image
+ * @typeParam W width, positive and integral
+ * @typeParam H height, positive and integral
+ */
+export declare class Image<D extends 1 | 3 | 4 = 1 | 3 | 4, W extends number = number, H extends number = number> {
+    readonly data: Readonly<Uint8Array>;
+    readonly width: W;
+    readonly height: H;
+    readonly depth: D;
+    constructor(data: Readonly<Uint8Array>, width: W, height: H, depth: D);
+}

package/dist/dataset/image.js

@@ -0,0 +1,21 @@
+/**
+ * Raw image with type level dimensions.
+ *
+ * @typeParam D depth of the image
+ * @typeParam W width, positive and integral
+ * @typeParam H height, positive and integral
+ */
+export class Image {
+    data;
+    width;
+    height;
+    depth;
+    constructor(data, width, height, depth) {
+        this.data = data;
+        this.width = width;
+        this.height = height;
+        this.depth = depth;
+        if (data.length != width * height * depth)
+            throw new Error("data isn't of excepted size");
+    }
+}

package/dist/dataset/index.d.ts

@@ -1,5 +1,3 @@
-export type { Dataset } from './dataset.js';
-export { DatasetBuilder } from './dataset_builder.js';
-export { ImageLoader, TabularLoader, DataLoader, TextLoader } from './data_loader/index.js';
-export type { DataSplit } from './data/index.js';
-export { Data, TabularData, ImageData, TextData, ImagePreprocessing, TabularPreprocessing, TextPreprocessing, IMAGE_PREPROCESSING, TABULAR_PREPROCESSING, TEXT_PREPROCESSING } from './data/index.js';
+export { Dataset } from "./dataset.js";
+export * from "./types.js";
+export { Data, TabularData, ImageData, TextData, ImagePreprocessing, TabularPreprocessing, TextPreprocessing, IMAGE_PREPROCESSING, TABULAR_PREPROCESSING, TEXT_PREPROCESSING, } from "./data/index.js";

package/dist/dataset/index.js

@@ -1,3 +1,3 @@
-export { DatasetBuilder } from './dataset_builder.js';
-export { ImageLoader, TabularLoader, DataLoader, TextLoader } from './data_loader/index.js';
-export { Data, TabularData, ImageData, TextData, ImagePreprocessing, TabularPreprocessing, TextPreprocessing, IMAGE_PREPROCESSING, TABULAR_PREPROCESSING, TEXT_PREPROCESSING } from './data/index.js';
+export { Dataset } from "./dataset.js";
+export * from "./types.js";
+export { Data, TabularData, ImageData, TextData, ImagePreprocessing, TabularPreprocessing, TextPreprocessing, IMAGE_PREPROCESSING, TABULAR_PREPROCESSING, TEXT_PREPROCESSING, } from "./data/index.js";

package/dist/dataset/types.d.ts

@@ -0,0 +1,4 @@
+import { Image } from "./image.js";
+export { Image };
+export type Tabular = Partial<Record<string, string>>;
+export type Text = string;

package/dist/dataset/types.js

@@ -0,0 +1,2 @@
+import { Image } from "./image.js";
+export { Image };

package/dist/index.d.ts

@@ -14,3 +14,7 @@ export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
 export * as async_iterator from "./utils/async_iterator.js";
+export { Dataset } from "./dataset/index.js";
+export * from "./dataset/types.js";
+export * from "./types.js";
+export * as processing from "./processing.js";

package/dist/index.js

@@ -14,3 +14,7 @@ export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
 export * as async_iterator from "./utils/async_iterator.js";
+export { Dataset } from "./dataset/index.js";
+export * from "./dataset/types.js"; // TODO merge with above
+export * from "./types.js";
+export * as processing from "./processing.js";

package/dist/models/gpt/model.js

@@ -30,6 +30,8 @@ class GPTModel extends tf.LayersModel {
         return this.config;
     }
     compile() {
+        if (this.optimizer !== undefined)
+            return;
         this.optimizer = this.config.weightDecay !== 0
             ? getCustomAdam(this, this.config.lr, this.config.weightDecay)
             : tf.train.adam(this.config.lr);

package/dist/models/model.d.ts

@@ -1,6 +1,5 @@
 import type tf from "@tensorflow/tfjs";
 import type { WeightsContainer } from "../index.js";
-import type { Dataset } from "../dataset/index.js";
 import type { BatchLogs, EpochLogs } from "./logs.js";
 export type Prediction = tf.Tensor;
 export type Sample = tf.Tensor;
@@ -23,7 +22,7 @@ export declare abstract class Model implements Disposable {
      * @param tracker watch the various steps
      * @yields on every epoch, training can be stop by `return`ing it
      */
-    abstract train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
+    abstract train(trainingData: tf.data.Dataset<tf.TensorContainer>, validationData?: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<BatchLogs, EpochLogs>;
     /** Predict likely values */
     abstract predict(input: Sample): Promise<Prediction>;
     /**