@stellarapp/tfjs-stellar 1.0.3 → 1.0.5

package/dist/src/models/gpt_model.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"gpt_model.d.ts","sourceRoot":"","sources":["../../../src/models/gpt_model.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,cAAc,EAAE,MAAM,cAAc,CAAC;AACnD,OAAO,EAAE,QAAQ,EAAE,KAAK,YAAY,EAAE,MAAM,oBAAoB,CAAC;AACjE,OAAO,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAC9C,OAAO,EAAE,KAAK,aAAa,EAAE,MAAM,8CAA8C,CAAC;AAIlF,MAAM,WAAW,YAAa,SAAQ,YAAY;IAC9C;;OAEG;IACH,QAAQ,EAAE,MAAM,CAAC;IACjB;;OAEG;IACH,SAAS,EAAE,MAAM,CAAC;IAClB;;OAEG;IACH,QAAQ,EAAE,MAAM,CAAC;IACjB;;;OAGG;IACH,SAAS,EAAE,MAAM,CAAC;IAClB;;;;;;OAMG;IACH,iBAAiB,CAAC,EAAE,OAAO,CAAC;CAC/B;AAGD;;;;;;;;;;;;;;;;;;;;;;GAsBG;AACH,qBAAa,QAAS,SAAQ,QAAQ;IAClC,MAAM,CAAC,SAAS,SAAc;IAE9B,SAAS,CAAC,QAAQ,CAAC,QAAQ,EAAE,MAAM,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,SAAS,EAAE,MAAM,CAAC;IACrC,SAAS,CAAC,QAAQ,CAAC,QAAQ,EAAE,MAAM,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,SAAS,EAAE,MAAM,CAAC;IACrC,SAAS,CAAC,QAAQ,CAAC,iBAAiB,EAAE,OAAO,CAAC;IAI9C,SAAS,CAAC,QAAQ,CAAC,eAAe,EAAE,MAAM,CAAC;IAG3C,SAAS,CAAC,kBAAkB,CAAC,EAAE,EAAE,CAAC,QAAQ,CAAC;IAG3C;;;OAGG;gBACS,IAAI,EAAE,YAAY;cAgBX,QAAQ,CACvB,EAAE,EAAE,EAAE,CAAC,MAAM,EACb,EAAE,EAAE,EAAE,CAAC,MAAM,EACb,SAAS,EAAE,EAAE,CAAC,MAAM,GAAG,SAAS,EAChC,aAAa,EAAE,cAAc,EAC7B,WAAW,CAAC,EAAE;QAAE,CAAC,GAAG,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM,GAAG,SAAS,CAAA;KAAE;;;;IAsC1D;;;;;OAKG;IACM,gBAAgB,CAAC,KAAK,EAAE,EAAE,CAAC,QAAQ,EAAE,QAAQ,EAAE,gBAAgB,GAAG,EAAE,CAAC,QAAQ;IA4B7E,KAAK,CAAC,UAAU,CAAC,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,IAAI;IA+B/C,OAAO,IAAI,aAAa;IAMxB,SAAS;;;;;;;;CAiBrB"}

package/dist/src/models/gpt_model.js

@@ -1,154 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-import { LlmModel } from "@/models/llm_model";
-import { GPT2DecoderBlock } from "@/layers/gpt_decoder_block";
-/**
- * This is a subclass of tf.Sequential that creating a GPT-like model and
- * automatically handles padding (and masking) the vocab size for hardware
- * efficiency.
- *
- * Example:
- *
- * ```javascript
- *
- * const model = new GptModel({ numLayers: 1, numHeads: 1, embedDim: 16, vocabSize: 64 });
- * model.compile({ loss: "sparseCategoricalCrossentropy", optimizer: "adam" });
- *
- * // use fitDataset() instead of fit for masking support
- * model.fitDataset(your_batched_generator_dataset, { epochs: 1 });
- *
- * const kv_cache = new KvCacheContainer(your_preferred_max_sequence_length);
- *
- * // use generate() and predictNextToken() instead of predict() for masking and auto memory cleanup
- * model.generate(tokenized_tensor1d_input, kv_cache, onPredict_callback)
- *
- *
- * ```
- */
-export class GptModel extends LlmModel {
-    static className = "GptModel";
-    numHeads;
-    numLayers;
-    embedDim;
-    vocabSize;
-    padToMultipleOf64;
-    // this is kept for reproducibility and model history but is not important since
-    // it can be calculated mathematically
-    vocabSizePadded;
-    // the amount to pad the embedding vocab size and dense output units count
-    vocab_padding_mask;
-    /**
-     * DO NOT add layers in the constructor or it will break tf.loadLayersModel().
-     * It should be done in build() instead.
-     */
-    constructor(args) {
-        const { numHeads, numLayers, embedDim, vocabSize, padToMultipleOf64 = true, ...rest } = args;
-        super({ name: "model", ...rest });
-        this.numHeads = numHeads;
-        this.numLayers = numLayers;
-        this.embedDim = embedDim;
-        this.vocabSize = vocabSize;
-        this.padToMultipleOf64 = padToMultipleOf64;
-        this.vocabSizePadded = this.padToMultipleOf64
-            ? Math.ceil(this.vocabSize / 64) * 64
-            : this.vocabSize;
-    }
-    fitBatch(xs, ys, loss_mask, loss_function, other_masks) {
-        let y_pred;
-        // forward pass, calculate loss
-        const { value: loss, grads } = tf.variableGrads(() => {
-            y_pred = this.apply(xs, {
-                training: true,
-                ...other_masks
-            });
-            // apply vocab pad masking
-            if (this.vocab_padding_mask) {
-                y_pred = y_pred.add(this.vocab_padding_mask);
-            }
-            y_pred = tf.softmax(y_pred);
-            // manually dispose later instead of the built-in disposal from variableGrads
-            tf.keep(y_pred);
-            const loss = loss_mask
-                ? loss_function(ys, y_pred).mul(loss_mask)
-                : loss_function(ys, y_pred);
-            return loss.mean();
-        });
-        // backpropagation
-        this.optimizer.applyGradients(grads);
-        return {
-            y_pred: y_pred,
-            loss
-        };
-    }
-    /**
-     * Overrides LlmModel.predictNextToken to add softmax before argMax because the final
-     * dense layer doesn't have an activation.
-     *
-     * TODO: implement temperature and multinomial sampling so that the model has varied outputs
-     */
-    predictNextToken(input, kv_cache) {
-        if (input.shape[0] != 1) {
-            throw Error(`GptModel.predictNextToken: ${this.name} expects an input with a batch size of 1`);
-        }
-        return tf.tidy(() => {
-            // comes back as [batch, sequence_length, vocab_size]
-            const prediction = this.apply(input, { kvCache: kv_cache });
-            const [batch_size, sequence_length, vocab_size] = prediction.shape;
-            // get the last token
-            const next_token = this.vocab_padding_mask != undefined
-                ? prediction
-                    .slice([0, sequence_length - 1, 0], [batch_size, 1, vocab_size])
-                    .add(this.vocab_padding_mask)
-                    .softmax()
-                    .argMax(2)
-                : prediction
-                    .slice([0, sequence_length - 1, 0], [batch_size, 1, vocab_size])
-                    .softmax()
-                    .argMax(2);
-            return next_token;
-        });
-    }
-    build(inputShape) {
-        const actual_vocab_size = this.vocabSizePadded
-            ? this.vocabSizePadded
-            : this.padToMultipleOf64
-                ? Math.ceil(this.vocabSize / 64) * 64
-                : this.vocabSize;
-        if (this.layers.length == 0) {
-            [
-                tf.layers.embedding({ inputDim: actual_vocab_size, outputDim: this.embedDim, batchInputShape: [null, null] }),
-                ...Array(this.numLayers).fill(0).map(_ => new GPT2DecoderBlock({ numHeads: this.numHeads, embedDim: this.embedDim })),
-                tf.layers.dense({ units: actual_vocab_size })
-            ].forEach(layer => this.add(layer));
-        }
-        if (this.vocab_padding_mask) {
-            this.vocab_padding_mask.dispose();
-        }
-        if (this.padToMultipleOf64 && actual_vocab_size > this.vocabSize) {
-            this.vocab_padding_mask = tf.tidy(() => tf.where(
-            // Create a mask of padded vocab length, values after the index "vocabSize"
-            // are set to -1e7 to mask out those positions so that softmax will ignore
-            // them. This mask is added to the final dense layer's output
-            tf.range(0, actual_vocab_size).greaterEqual(this.vocabSize), -1e7, 0).toFloat());
-        }
-        super.build(inputShape);
-    }
-    dispose() {
-        this.vocab_padding_mask?.dispose();
-        return super.dispose();
-    }
-    getConfig() {
-        const base_config = super.getConfig();
-        const config = {
-            numHeads: this.numHeads,
-            numLayers: this.numLayers,
-            embedDim: this.embedDim,
-            vocabSize: this.vocabSize,
-            vocabSizePadded: this.vocabSizePadded,
-            padToMultipleOf64: this.padToMultipleOf64
-        };
-        Object.assign(config, base_config);
-        return config;
-    }
-}
-tf.serialization.registerClass(GptModel);
-//# sourceMappingURL=gpt_model.js.map

package/dist/src/models/gpt_model.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"gpt_model.js","sourceRoot":"","sources":["../../../src/models/gpt_model.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,QAAQ,EAAqB,MAAM,oBAAoB,CAAC;AAGjE,OAAO,EAAE,gBAAgB,EAAE,MAAM,4BAA4B,CAAC;AAgC9D;;;;;;;;;;;;;;;;;;;;;;GAsBG;AACH,MAAM,OAAO,QAAS,SAAQ,QAAQ;IAClC,MAAM,CAAC,SAAS,GAAG,UAAU,CAAC;IAEX,QAAQ,CAAS;IACjB,SAAS,CAAS;IAClB,QAAQ,CAAS;IACjB,SAAS,CAAS;IAClB,iBAAiB,CAAU;IAE9C,gFAAgF;IAChF,sCAAsC;IACnB,eAAe,CAAS;IAE3C,0EAA0E;IAChE,kBAAkB,CAAe;IAG3C;;;OAGG;IACH,YAAY,IAAkB;QAC1B,MAAM,EAAE,QAAQ,EAAE,SAAS,EAAE,QAAQ,EAAE,SAAS,EAAE,iBAAiB,GAAG,IAAI,EAAE,GAAG,IAAI,EAAE,GAAG,IAAI,CAAC;QAE7F,KAAK,CAAC,EAAE,IAAI,EAAE,OAAO,EAAE,GAAG,IAAI,EAAE,CAAC,CAAC;QAElC,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,SAAS,GAAG,SAAS,CAAC;QAC3B,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,SAAS,GAAG,SAAS,CAAC;QAC3B,IAAI,CAAC,iBAAiB,GAAG,iBAAiB,CAAC;QAC3C,IAAI,CAAC,eAAe,GAAG,IAAI,CAAC,iBAAiB;YACzC,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,SAAS,GAAG,EAAE,CAAC,GAAG,EAAE;YACrC,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC;IACzB,CAAC;IAGkB,QAAQ,CACvB,EAAa,EACb,EAAa,EACb,SAAgC,EAChC,aAA6B,EAC7B,WAAsD;QAEtD,IAAI,MAAiB,CAAC;QAEtB,+BAA+B;QAC/B,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,CAAC,aAAa,CAAC,GAAG,EAAE;YACjD,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,EAAE,EAAE;gBACpB,QAAQ,EAAE,IAAI;gBACd,GAAG,WAAW;aACjB,CAAc,CAAC;YAEhB,0BAA0B;YAC1B,IAAI,IAAI,CAAC,kBAAkB,EAAE,CAAC;gBAC1B,MAAM,GAAG,MAAM,CAAC,GAAG,CAAC,IAAI,CAAC,kBAAkB,CAAC,CAAC;YACjD,CAAC;YAED,MAAM,GAAG,EAAE,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC;YAE5B,6EAA6E;YAC7E,EAAE,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;YAEhB,MAAM,IAAI,GAAG,SAAS;gBAClB,CAAC,CAAC,aAAa,CAAC,EAAE,EAAE,MAAM,CAAC,CAAC,GAAG,CAAC,SAAS,CAAC;gBAC1C,CAAC,CAAC,aAAa,CAAC,EAAE,EAAE,MAAM,CAAC,CAAC;YAEhC,OAAO,IAAI,CAAC,IAAI,EAAe,CAAC;QACpC,CAAC,CAAC,CAAC;QAEH,kBAAkB;QAClB,IAAI,CAAC,SAAS,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC;QAErC,OAAO;YACH,MAAM,EAAE,MAAO;YACf,IAAI;SACP,CAAC;IACN,CAAC;IAGD;;;;;OAKG;IACM,gBAAgB,CAAC,KAAkB,EAAE,QAA0B;QACpE,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,IAAI,CAAC,EAAE,CAAC;YACtB,MAAM,KAAK,CAAC,8BAA8B,IAAI,CAAC,IAAI,0CAA0C,CAAC,CAAC;QACnG,CAAC;QAED,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,qDAAqD;YACrD,MAAM,UAAU,GAAG,IAAI,CAAC,KAAK,CAAC,KAAK,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAc,CAAC;YAEzE,MAAM,CAAC,UAAU,EAAE,eAAe,EAAE,UAAU,CAAC,GAAG,UAAU,CAAC,KAAK,CAAC;YAEnE,qBAAqB;YACrB,MAAM,UAAU,GAAG,IAAI,CAAC,kBAAkB,IAAI,SAAS;gBACnD,CAAC,CAAC,UAAU;qBACP,KAAK,CAAC,CAAC,CAAC,EAAE,eAAe,GAAG,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,UAAU,EAAE,CAAC,EAAE,UAAU,CAAC,CAAC;qBAC/D,GAAG,CAAC,IAAI,CAAC,kBAAkB,CAAC;qBAC5B,OAAO,EAAE;qBACT,MAAM,CAAC,CAAC,CAAC;gBACd,CAAC,CAAC,UAAU;qBACP,KAAK,CAAC,CAAC,CAAC,EAAE,eAAe,GAAG,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,UAAU,EAAE,CAAC,EAAE,UAAU,CAAC,CAAC;qBAC/D,OAAO,EAAE;qBACT,MAAM,CAAC,CAAC,CAAC,CAAC;YAEnB,OAAO,UAAyB,CAAC;QACrC,CAAC,CAAC,CAAA;IACN,CAAC;IAGQ,KAAK,CAAC,UAAkC;QAC7C,MAAM,iBAAiB,GAAG,IAAI,CAAC,eAAe;YAC1C,CAAC,CAAC,IAAI,CAAC,eAAe;YACtB,CAAC,CAAC,IAAI,CAAC,iBAAiB;gBACpB,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,SAAS,GAAG,EAAE,CAAC,GAAG,EAAE;gBACrC,CAAC,CAAC,IAAI,CAAC,SAAS,CAAA;QAExB,IAAI,IAAI,CAAC,MAAM,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YAC1B;gBACI,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,EAAE,QAAQ,EAAE,iBAAiB,EAAE,SAAS,EAAE,IAAI,CAAC,QAAQ,EAAE,eAAe,EAAE,CAAC,IAAI,EAAE,IAAI,CAAC,EAAE,CAAC;gBAC7G,GAAG,KAAK,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,IAAI,gBAAgB,CAAC,EAAE,QAAQ,EAAE,IAAI,CAAC,QAAQ,EAAE,QAAQ,EAAE,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAC;gBACrH,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,EAAE,KAAK,EAAE,iBAAiB,EAAE,CAAC;aAChD,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE,CAAC,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC,CAAA;QACvC,CAAC;QAED,IAAI,IAAI,CAAC,kBAAkB,EAAE,CAAC;YAC1B,IAAI,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC;QACtC,CAAC;QAED,IAAI,IAAI,CAAC,iBAAiB,IAAI,iBAAiB,GAAG,IAAI,CAAC,SAAS,EAAE,CAAC;YAC/D,IAAI,CAAC,kBAAkB,GAAG,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,KAAK;YAC5C,2EAA2E;YAC3E,0EAA0E;YAC1E,6DAA6D;YAC7D,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,iBAAiB,CAAC,CAAC,YAAY,CAAC,IAAI,CAAC,SAAS,CAAC,EAAE,CAAC,GAAG,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC,CAAA;QACxF,CAAC;QAED,KAAK,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;IAC5B,CAAC;IAGQ,OAAO;QACZ,IAAI,CAAC,kBAAkB,EAAE,OAAO,EAAE,CAAC;QACnC,OAAO,KAAK,CAAC,OAAO,EAAE,CAAC;IAC3B,CAAC;IAGQ,SAAS;QACd,MAAM,WAAW,GAAG,KAAK,CAAC,SAAS,EAAE,CAAC;QAEtC,MAAM,MAAM,GAAG;YACX,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,SAAS,EAAE,IAAI,CAAC,SAAS;YACzB,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,SAAS,EAAE,IAAI,CAAC,SAAS;YACzB,eAAe,EAAE,IAAI,CAAC,eAAe;YACrC,iBAAiB,EAAE,IAAI,CAAC,iBAAiB;SAC5C,CAAA;QAED,MAAM,CAAC,MAAM,CAAC,MAAM,EAAE,WAAW,CAAC,CAAC;QAEnC,OAAO,MAAM,CAAC;IAClB,CAAC;;AAKL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,QAAQ,CAAC,CAAC"}

package/dist/src/models/index.d.ts

@@ -1,3 +0,0 @@
-export * from "@/models/gpt_model";
-export * from "@/models/u_net";
-//# sourceMappingURL=index.d.ts.map

package/dist/src/models/index.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../src/models/index.ts"],"names":[],"mappings":"AAAA,cAAc,oBAAoB,CAAC;AACnC,cAAc,gBAAgB,CAAC"}

package/dist/src/models/index.js

@@ -1,3 +0,0 @@
-export * from "@/models/gpt_model";
-export * from "@/models/u_net";
-//# sourceMappingURL=index.js.map

package/dist/src/models/index.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../../../src/models/index.ts"],"names":[],"mappings":"AAAA,cAAc,oBAAoB,CAAC;AACnC,cAAc,gBAAgB,CAAC"}

package/dist/src/models/llm_model.d.ts

@@ -1,87 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-import { Dataset, type LossOrMetricFn } from "@/tfjs_types";
-import { KvCacheContainer } from "@/kv_cache";
-export interface LlmModelArgs extends tf.SequentialArgs {
-}
-interface DatasetArgs extends tf.TensorContainerObject {
-    xs: tf.Tensor;
-    ys: tf.Tensor;
-    loss_mask?: tf.Tensor;
-    packing_mask?: tf.Tensor;
-}
-/**
- * This class overrides the `fitDataset()` function of tf.Sequential to support loss
- * and packing masking. Use the `generate()` function to autoregressively predict the
- * next, set `stopPredicting=true` to stop.
- */
-export declare class LlmModel extends tf.Sequential {
-    static className: string;
-    private stopPredicting_;
-    constructor(args: LlmModelArgs);
-    /**
-     * Returns the metric functions and names so that metrics can be reported
-     * as they are in the base version of model.fitDataset
-     *
-     * e.g. "categoricalAccuracy" should be reported as "acc"
-     */
-    protected getMetricFunctions(): {
-        metric_fn: import("@tensorflow/tfjs-layers/dist/types").LossOrMetricFn;
-        metric_label: string;
-    }[];
-    /**
-     * Get exactly one loss function from the loss function provided in `model.compile()`.
-     * If a string identifier was used, convert it to the actual loss function.
-     */
-    protected getLossFunction(): LossOrMetricFn;
-    /**
-     * Train on a `tf.data.generator` dataset. See https://js.tensorflow.org/api/latest/#data.generator.
-     *
-     * The generator should yield `xs`, `ys`, `loss_mask` (if fine-tuning), and
-     * `packing_mask` (if sequence packing was done)
-     *
-     * @param tfdataset an instance of a `tf.Dataset` generator
-     * @param args a ModelFitDatasetArgs
-     */
-    fitDataset<T = DatasetArgs>(tfdataset: Dataset<T>, args: tf.ModelFitDatasetArgs<T>): Promise<any>;
-    /**
-     * Run the core forward and backward propagation on one training batch. This
-     * should be called within a `tf.tidy()`.
-     *
-     * @param xs the sample/input tensor
-     * @param ys the label/target tensor
-     * @param loss_mask a loss mask to ignore the prediction's non-assistant tokens
-     * @param loss_function the model's loss function
-     * @param other_masks other masks used by the model's layers e.g. packing mask, causal mask
-     */
-    protected fitBatch(xs: tf.Tensor, ys: tf.Tensor, loss_mask: tf.Tensor | undefined, loss_function: LossOrMetricFn, other_masks?: {
-        [key: string]: tf.Tensor | undefined;
-    }): {
-        y_pred: tf.Tensor<tf.Rank>;
-        loss: tf.Scalar;
-    };
-    compile(args: tf.ModelCompileArgs): void;
-    /**
-     * Autoregressively generate the next token until `model.stopPredicting` is set
-     * to `true` or the KV cache reaches its maximum sequence length. For a single chat
-     * session, the input should only be the most recent prompt(s). The KV cache stores
-     * the prior chat history up until the most recent chat.
-     *
-     * @param input tokenized input of the newest chat
-     * @param kv_cache an instance of a KV cache container
-     * @param onPredict callback function to receive the most recent token predicted
-     */
-    generate(input: tf.Tensor1D, kv_cache: KvCacheContainer, onPredict: (token: tf.Tensor) => Promise<void>): Promise<void>;
-    /**
-     * Given a tokenized sentence, predict the next token (word).
-     * A normal prediction is ran to get an output with the shape
-     * `[ batch_size, sentence_length, vocab_size ]` and the `vocab_size`
-     * position with the highest scored probability in the last
-     * position of `sentence_length` is returned as the next predicted
-     * token.
-     */
-    predictNextToken(input: tf.Tensor2D, kv_cache: KvCacheContainer): tf.Tensor2D;
-    get stopPredicting(): boolean;
-    set stopPredicting(stop: boolean);
-}
-export {};
-//# sourceMappingURL=llm_model.d.ts.map

package/dist/src/models/llm_model.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"llm_model.d.ts","sourceRoot":"","sources":["../../../src/models/llm_model.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,OAAO,EAAE,KAAK,cAAc,EAAE,MAAM,cAAc,CAAC;AAE5D,OAAO,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAK9C,MAAM,WAAW,YAAa,SAAQ,EAAE,CAAC,cAAc;CACtD;AAGD,UAAU,WAAY,SAAQ,EAAE,CAAC,qBAAqB;IAClD,EAAE,EAAE,EAAE,CAAC,MAAM,CAAC;IACd,EAAE,EAAE,EAAE,CAAC,MAAM,CAAC;IACd,SAAS,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC;IACtB,YAAY,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC;CAC5B;AAGD;;;;GAIG;AACH,qBAAa,QAAS,SAAQ,EAAE,CAAC,UAAU;IACvC,MAAM,CAAC,SAAS,SAAc;IAE9B,OAAO,CAAC,eAAe,CAAiB;gBAE5B,IAAI,EAAE,YAAY;IAM9B;;;;;OAKG;IACH,SAAS,CAAC,kBAAkB;;;;IAU5B;;;OAGG;IACH,SAAS,CAAC,eAAe,IAAI,cAAc;IAoC3C;;;;;;;;OAQG;IACY,UAAU,CAAC,CAAC,GAAG,WAAW,EAAE,SAAS,EAAE,OAAO,CAAC,CAAC,CAAC,EAAE,IAAI,EAAE,EAAE,CAAC,mBAAmB,CAAC,CAAC,CAAC,GAAG,OAAO,CAAC,GAAG,CAAC;IAqHhH;;;;;;;;;OASG;IACH,SAAS,CAAC,QAAQ,CACd,EAAE,EAAE,EAAE,CAAC,MAAM,EACb,EAAE,EAAE,EAAE,CAAC,MAAM,EACb,SAAS,EAAE,EAAE,CAAC,MAAM,GAAG,SAAS,EAChC,aAAa,EAAE,cAAc,EAC7B,WAAW,CAAC,EAAE;QAAE,CAAC,GAAG,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM,GAAG,SAAS,CAAA;KAAE,GACvD;QACC,MAAM,EAAE,EAAE,CAAC,MAAM,CAAC,EAAE,CAAC,IAAI,CAAC,CAAC;QAC3B,IAAI,EAAE,EAAE,CAAC,MAAM,CAAC;KACnB;IA+BQ,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,gBAAgB,GAAG,IAAI;IASjD;;;;;;;;;OASG;IACU,QAAQ,CAAC,KAAK,EAAE,EAAE,CAAC,QAAQ,EAAE,QAAQ,EAAE,gBAAgB,EAAE,SAAS,EAAE,CAAC,KAAK,EAAE,EAAE,CAAC,MAAM,KAAK,OAAO,CAAC,IAAI,CAAC;IA2BpH;;;;;;;OAOG;IACI,gBAAgB,CAAC,KAAK,EAAE,EAAE,CAAC,QAAQ,EAAE,QAAQ,EAAE,gBAAgB;IAmBtE,IAAI,cAAc,IAKO,OAAO,CAH/B;IAGD,IAAI,cAAc,CAAC,IAAI,EAAE,OAAO,EAE/B;CAEJ"}

package/dist/src/models/llm_model.js

@@ -1,245 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-import { sparseCategoricalCrossentropy } from "@tensorflow/tfjs-layers/dist/losses";
-import { causal as generateCausalMask } from "@/masks";
-import * as losses from "@/losses";
-;
-/**
- * This class overrides the `fitDataset()` function of tf.Sequential to support loss
- * and packing masking. Use the `generate()` function to autoregressively predict the
- * next, set `stopPredicting=true` to stop.
- */
-export class LlmModel extends tf.Sequential {
-    static className = "LlmModel";
-    stopPredicting_ = true;
-    constructor(args) {
-        args.name = args.name ?? "model";
-        super(args);
-    }
-    /**
-     * Returns the metric functions and names so that metrics can be reported
-     * as they are in the base version of model.fitDataset
-     *
-     * e.g. "categoricalAccuracy" should be reported as "acc"
-     */
-    getMetricFunctions() {
-        const [loss, ...metric_fn_names] = this.metricsNames;
-        return this.metricsTensors.map((metric_tensor, index) => ({
-            metric_fn: metric_tensor[0],
-            metric_label: metric_fn_names[index]
-        }));
-    }
-    /**
-     * Get exactly one loss function from the loss function provided in `model.compile()`.
-     * If a string identifier was used, convert it to the actual loss function.
-     */
-    getLossFunction() {
-        let loss = this.loss;
-        if (Array.isArray(loss)) {
-            loss = loss[0];
-        }
-        if (typeof loss == "string") {
-            if (loss == "sparseCategoricalCrossentropy") {
-                return sparseCategoricalCrossentropy;
-                /* throw Error("LlmModel.getLossFunction: TFJS's sparseCategoricalCrossentropy" +
-                    " is not truly sparse, it simply converts it to onehot." +
-                    " Use categoricalCrossentropy instead. See" +
-                    " https://github.com/tensorflow/tfjs/blob/0fc04d958ea592f3b8db79a8b3b497b5c8904097/tfjs-layers/src/losses.ts#L143-L146"); */
-            }
-            const loss_id = loss;
-            const loss_fn = (losses[loss_id] ??
-                tf.losses[loss_id] ??
-                tf.metrics[loss_id]);
-            if (loss_fn) {
-                return loss_fn;
-            }
-            else {
-                throw Error(`LlmModel.getLossFunction: ${loss_id} is not a valid loss function`);
-            }
-        }
-        else if (typeof loss == "function") {
-            return loss;
-        }
-        throw Error("LlmModel.getLossFunction: the loss function's type should be string or function");
-    }
-    /**
-     * Train on a `tf.data.generator` dataset. See https://js.tensorflow.org/api/latest/#data.generator.
-     *
-     * The generator should yield `xs`, `ys`, `loss_mask` (if fine-tuning), and
-     * `packing_mask` (if sequence packing was done)
-     *
-     * @param tfdataset an instance of a `tf.Dataset` generator
-     * @param args a ModelFitDatasetArgs
-     */
-    async fitDataset(tfdataset, args) {
-        this.stopTraining = false;
-        const dataset = tfdataset;
-        const { epochs, callbacks } = args;
-        const metric_functions = this.getMetricFunctions();
-        const loss_function = this.getLossFunction();
-        this.lossFunctions = [loss_function];
-        const { onBatchBegin, onBatchEnd, onEpochBegin, onEpochEnd, onTrainBegin, onTrainEnd, } = callbacks ?? {};
-        await onTrainBegin?.();
-        let cached_causal_mask = undefined;
-        for (let epoch = 0; epoch < epochs; epoch++) {
-            await onEpochBegin?.(epoch);
-            let batch = 0;
-            let total_samples = 0;
-            const accumulated_epoch_metrics = {};
-            // loop through dataset using its iterator
-            const iterator = await dataset.iterator();
-            let sample = await iterator.next();
-            while (!sample.done) {
-                const batch_metrics = { batch };
-                const { xs, ys, loss_mask, packing_mask } = sample.value;
-                const batch_size = xs.shape[0];
-                total_samples += batch_size; // for epoch metrics averaging
-                if (xs.shape.length != 2) {
-                    throw Error(`LlmModel.fitDataset: ${this.name} the generator dataset should be batched, run: dataset.batch(batch_size)`);
-                }
-                // pre-calculate the causal attention mask and reuse it for all attention layers,
-                const seq_length = xs.shape[xs.shape.length - 1];
-                if (!cached_causal_mask || cached_causal_mask.shape[0] != seq_length) {
-                    cached_causal_mask = generateCausalMask(seq_length, seq_length);
-                }
-                await onBatchBegin?.(batch);
-                tf.tidy(() => {
-                    const { y_pred, loss } = this.fitBatch(xs, ys, loss_mask, loss_function, {
-                        packingMask: packing_mask,
-                        causalMask: cached_causal_mask
-                    });
-                    const loss_value = (loss.dataSync())[0];
-                    batch_metrics.loss = loss_value;
-                    accumulated_epoch_metrics.loss = (accumulated_epoch_metrics.loss || 0) + loss_value * batch_size;
-                    // calculate and store metrics
-                    for (const { metric_fn, metric_label } of metric_functions) {
-                        const metric_sum = metric_fn(ys, y_pred).mean();
-                        const metric_value = (metric_sum.dataSync())[0];
-                        batch_metrics[metric_label] = metric_value; // / batch_size;
-                        accumulated_epoch_metrics[metric_label] = (accumulated_epoch_metrics[metric_label] || 0) + metric_value * batch_size;
-                    }
-                    tf.dispose(y_pred);
-                });
-                tf.dispose(xs);
-                tf.dispose(ys);
-                tf.dispose(loss_mask);
-                if (packing_mask) {
-                    tf.dispose(packing_mask);
-                }
-                await onBatchEnd?.(batch, batch_metrics);
-                // so that stop training works
-                await tf.nextFrame();
-                if (this.stopTraining) {
-                    break;
-                }
-                sample = await iterator.next();
-                batch++;
-            }
-            for (const metric in accumulated_epoch_metrics) {
-                accumulated_epoch_metrics[metric] = accumulated_epoch_metrics[metric] / total_samples;
-            }
-            await onEpochEnd?.(epoch, accumulated_epoch_metrics);
-            if (this.stopTraining) {
-                break;
-            }
-        }
-        tf.dispose(cached_causal_mask);
-        await onTrainEnd?.();
-        return {};
-    }
-    /**
-     * Run the core forward and backward propagation on one training batch. This
-     * should be called within a `tf.tidy()`.
-     *
-     * @param xs the sample/input tensor
-     * @param ys the label/target tensor
-     * @param loss_mask a loss mask to ignore the prediction's non-assistant tokens
-     * @param loss_function the model's loss function
-     * @param other_masks other masks used by the model's layers e.g. packing mask, causal mask
-     */
-    fitBatch(xs, ys, loss_mask, loss_function, other_masks) {
-        let y_pred;
-        // forward pass, calculate loss
-        const { value: loss, grads } = tf.variableGrads(() => {
-            // prediction has shape [batch, sequence_length, vocab_size]
-            y_pred = this.apply(xs, {
-                training: true,
-                ...other_masks
-            });
-            // manually dispose later instead of the built-in disposal from variableGrads
-            tf.keep(y_pred);
-            const loss = loss_mask
-                ? loss_function(ys, y_pred).mul(loss_mask)
-                : loss_function(ys, y_pred);
-            return loss.mean();
-        });
-        // backpropagation
-        this.optimizer.applyGradients(grads);
-        return {
-            y_pred: y_pred,
-            loss
-        };
-    }
-    compile(args) {
-        if (args.loss == "categoricalCrossentropy") {
-            throw Error(`LlmModel.compile: use sparseCategoricalCrossentropy loss (along with onehot encoded labels) instead of categoricalCrossEntropy`);
-        }
-        super.compile(args);
-    }
-    /**
-     * Autoregressively generate the next token until `model.stopPredicting` is set
-     * to `true` or the KV cache reaches its maximum sequence length. For a single chat
-     * session, the input should only be the most recent prompt(s). The KV cache stores
-     * the prior chat history up until the most recent chat.
-     *
-     * @param input tokenized input of the newest chat
-     * @param kv_cache an instance of a KV cache container
-     * @param onPredict callback function to receive the most recent token predicted
-     */
-    async generate(input, kv_cache, onPredict) {
-        if (kv_cache.size >= kv_cache.maxSequenceLength) {
-            throw Error(`LlmModel.generate: ${this.name} KV cache's size reached the maxSequenceLength (${kv_cache.maxSequenceLength})`);
-        }
-        this.stopPredicting = false;
-        let current_token = tf.tidy(() => input.expandDims(0)); // it's 2D because of the required batch dimension
-        while (!this.stopPredicting && kv_cache.size < kv_cache.maxSequenceLength) {
-            // add a batch dimension because forward pass requires inputs batched
-            const next_token = tf.tidy(() => this.predictNextToken(current_token, kv_cache));
-            // pass back the predicted token, without the batch dim,
-            const unbatched_next_token = tf.tidy(() => next_token.squeeze([0]));
-            await onPredict(unbatched_next_token);
-            unbatched_next_token.dispose();
-            current_token.dispose();
-            current_token = next_token;
-        }
-        tf.dispose(current_token);
-    }
-    /**
-     * Given a tokenized sentence, predict the next token (word).
-     * A normal prediction is ran to get an output with the shape
-     * `[ batch_size, sentence_length, vocab_size ]` and the `vocab_size`
-     * position with the highest scored probability in the last
-     * position of `sentence_length` is returned as the next predicted
-     * token.
-     */
-    predictNextToken(input, kv_cache) {
-        if (input.shape[0] != 1) {
-            throw Error(`LlmModel.predictNextToken: ${this.name} expects an input with a batch size of 1`);
-        }
-        return tf.tidy(() => {
-            // comes back as [batch, sequence_length, vocab_size]
-            const prediction = this.apply(input, { kvCache: kv_cache });
-            const [batch_size, sequence_length, vocab_size] = prediction.shape;
-            // get the last token
-            const next_token = prediction.slice([0, sequence_length - 1, 0], [batch_size, 1, vocab_size]).argMax(2);
-            return next_token;
-        });
-    }
-    get stopPredicting() {
-        return this.stopPredicting_;
-    }
-    set stopPredicting(stop) {
-        this.stopPredicting_ = stop;
-    }
-}
-tf.serialization.registerClass(LlmModel);
-//# sourceMappingURL=llm_model.js.map

package/dist/src/models/llm_model.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"llm_model.js","sourceRoot":"","sources":["../../../src/models/llm_model.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,6BAA6B,EAAE,MAAM,qCAAqC,CAAC;AAEpF,OAAO,EAAE,MAAM,IAAI,kBAAkB,EAAE,MAAM,SAAS,CAAC;AAEvD,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AAKlC,CAAC;AAWF;;;;GAIG;AACH,MAAM,OAAO,QAAS,SAAQ,EAAE,CAAC,UAAU;IACvC,MAAM,CAAC,SAAS,GAAG,UAAU,CAAC;IAEtB,eAAe,GAAY,IAAI,CAAC;IAExC,YAAY,IAAkB;QAC1B,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC,IAAI,IAAI,OAAO,CAAC;QACjC,KAAK,CAAC,IAAI,CAAC,CAAC;IAChB,CAAC;IAGD;;;;;OAKG;IACO,kBAAkB;QACxB,MAAM,CAAC,IAAI,EAAE,GAAG,eAAe,CAAC,GAAG,IAAI,CAAC,YAAY,CAAC;QAErD,OAAO,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,CAAC,aAAa,EAAE,KAAK,EAAE,EAAE,CAAC,CAAC;YACtD,SAAS,EAAE,aAAa,CAAC,CAAC,CAAC;YAC3B,YAAY,EAAE,eAAe,CAAC,KAAK,CAAC;SACvC,CAAC,CAAC,CAAA;IACP,CAAC;IAGD;;;OAGG;IACO,eAAe;QACrB,IAAI,IAAI,GAAG,IAAI,CAAC,IAAI,CAAC;QAErB,IAAI,KAAK,CAAC,OAAO,CAAC,IAAI,CAAC,EAAE,CAAC;YACtB,IAAI,GAAG,IAAI,CAAC,CAAC,CAAC,CAAC;QACnB,CAAC;QAED,IAAI,OAAO,IAAI,IAAI,QAAQ,EAAE,CAAC;YAC1B,IAAI,IAAI,IAAI,+BAA+B,EAAE,CAAC;gBAC1C,OAAO,6BAA6B,CAAC;gBACrC;;;gJAGgI;YACpI,CAAC;YAED,MAAM,OAAO,GAAG,IAAc,CAAC;YAE/B,MAAM,OAAO,GACT,CAAE,MAA8B,CAAC,OAAO,CAAC;gBACpC,EAAE,CAAC,MAA8B,CAAC,OAAO,CAAC;gBAC1C,EAAE,CAAC,OAA+B,CAAC,OAAO,CAAC,CAAmB,CAAA;YAEvE,IAAI,OAAO,EAAE,CAAC;gBACV,OAAO,OAAO,CAAA;YAClB,CAAC;iBAAM,CAAC;gBACJ,MAAM,KAAK,CAAC,6BAA6B,OAAO,+BAA+B,CAAC,CAAC;YACrF,CAAC;QACL,CAAC;aAAM,IAAI,OAAO,IAAI,IAAI,UAAU,EAAE,CAAC;YACnC,OAAO,IAAI,CAAC;QAChB,CAAC;QAED,MAAM,KAAK,CAAC,iFAAiF,CAAC,CAAC;IACnG,CAAC;IAGD;;;;;;;;OAQG;IACM,KAAK,CAAC,UAAU,CAAkB,SAAqB,EAAE,IAA+B;QAC7F,IAAI,CAAC,YAAY,GAAG,KAAK,CAAC;QAE1B,MAAM,OAAO,GAAG,SAAyC,CAAC;QAC1D,MAAM,EAAE,MAAM,EAAE,SAAS,EAAE,GAAG,IAAI,CAAC;QAEnC,MAAM,gBAAgB,GAAG,IAAI,CAAC,kBAAkB,EAAE,CAAC;QACnD,MAAM,aAAa,GAAG,IAAI,CAAC,eAAe,EAAE,CAAC;QAC7C,IAAI,CAAC,aAAa,GAAG,CAAC,aAAa,CAAC,CAAC;QAErC,MAAM,EACF,YAAY,EACZ,UAAU,EACV,YAAY,EACZ,UAAU,EACV,YAAY,EACZ,UAAU,GACb,GAAG,SAAkC,IAAI,EAAE,CAAC;QAE7C,MAAM,YAAY,EAAE,EAAE,CAAC;QAEvB,IAAI,kBAAkB,GAA0B,SAAS,CAAC;QAE1D,KAAK,IAAI,KAAK,GAAG,CAAC,EAAE,KAAK,GAAG,MAAM,EAAE,KAAK,EAAE,EAAE,CAAC;YAC1C,MAAM,YAAY,EAAE,CAAC,KAAK,CAAC,CAAC;YAE5B,IAAI,KAAK,GAAG,CAAC,CAAC;YACd,IAAI,aAAa,GAAG,CAAC,CAAC;YACtB,MAAM,yBAAyB,GAAiC,EAAE,CAAC;YAEnE,0CAA0C;YAC1C,MAAM,QAAQ,GAAG,MAAM,OAAO,CAAC,QAAQ,EAAE,CAAC;YAC1C,IAAI,MAAM,GAAG,MAAM,QAAQ,CAAC,IAAI,EAAE,CAAC;YAEnC,OAAO,CAAC,MAAM,CAAC,IAAI,EAAE,CAAC;gBAClB,MAAM,aAAa,GAAiC,EAAE,KAAK,EAAE,CAAC;gBAE9D,MAAM,EAAE,EAAE,EAAE,EAAE,EAAE,SAAS,EAAE,YAAY,EAAE,GAAG,MAAM,CAAC,KAAK,CAAC;gBACzD,MAAM,UAAU,GAAG,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;gBAC/B,aAAa,IAAI,UAAU,CAAC,CAAC,8BAA8B;gBAE3D,IAAI,EAAE,CAAC,KAAK,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;oBACvB,MAAM,KAAK,CAAC,wBAAwB,IAAI,CAAC,IAAI,0EAA0E,CAAC,CAAC;gBAC7H,CAAC;gBAED,iFAAiF;gBACjF,MAAM,UAAU,GAAG,EAAE,CAAC,KAAK,CAAC,EAAE,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC;gBAEjD,IAAI,CAAC,kBAAkB,IAAI,kBAAkB,CAAC,KAAK,CAAC,CAAC,CAAC,IAAI,UAAU,EAAE,CAAC;oBACnE,kBAAkB,GAAG,kBAAkB,CAAC,UAAU,EAAE,UAAU,CAAC,CAAC;gBACpE,CAAC;gBAED,MAAM,YAAY,EAAE,CAAC,KAAK,CAAC,CAAC;gBAE5B,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;oBACT,MAAM,EAAE,MAAM,EAAE,IAAI,EAAE,GAAG,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,EAAE,EAAE,SAAS,EAAE,aAAa,EAAE;wBACrE,WAAW,EAAE,YAAY;wBACzB,UAAU,EAAE,kBAAkB;qBACjC,CAAC,CAAA;oBAEF,MAAM,UAAU,GAAG,CAAC,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;oBAExC,aAAa,CAAC,IAAI,GAAG,UAAU,CAAC;oBAChC,yBAAyB,CAAC,IAAI,GAAG,CAAC,yBAAyB,CAAC,IAAI,IAAI,CAAC,CAAC,GAAG,UAAU,GAAG,UAAU,CAAC;oBAEjG,8BAA8B;oBAC9B,KAAK,MAAM,EAAE,SAAS,EAAE,YAAY,EAAE,IAAI,gBAAgB,EAAE,CAAC;wBACzD,MAAM,UAAU,GAAG,SAAS,CAAC,EAAE,EAAE,MAAO,CAAC,CAAC,IAAI,EAAE,CAAC;wBAEjD,MAAM,YAAY,GAAG,CAAC,UAAU,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;wBAEhD,aAAa,CAAC,YAAY,CAAC,GAAG,YAAY,CAAA,CAAA,gBAAgB;wBAC1D,yBAAyB,CAAC,YAAY,CAAC,GAAG,CAAC,yBAAyB,CAAC,YAAY,CAAC,IAAI,CAAC,CAAC,GAAG,YAAY,GAAG,UAAU,CAAC;oBACzH,CAAC;oBAED,EAAE,CAAC,OAAO,CAAC,MAAO,CAAC,CAAC;gBACxB,CAAC,CAAC,CAAA;gBAEF,EAAE,CAAC,OAAO,CAAC,EAAE,CAAC,CAAC;gBACf,EAAE,CAAC,OAAO,CAAC,EAAE,CAAC,CAAC;gBACf,EAAE,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC;gBAEtB,IAAI,YAAY,EAAE,CAAC;oBACf,EAAE,CAAC,OAAO,CAAC,YAAY,CAAC,CAAC;gBAC7B,CAAC;gBAED,MAAM,UAAU,EAAE,CAAC,KAAK,EAAE,aAAa,CAAC,CAAC;gBAEzC,8BAA8B;gBAC9B,MAAM,EAAE,CAAC,SAAS,EAAE,CAAC;gBAErB,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;oBACpB,MAAM;gBACV,CAAC;gBAED,MAAM,GAAG,MAAM,QAAQ,CAAC,IAAI,EAAE,CAAC;gBAC/B,KAAK,EAAE,CAAC;YACZ,CAAC;YAED,KAAK,MAAM,MAAM,IAAI,yBAAyB,EAAE,CAAC;gBAC7C,yBAAyB,CAAC,MAAM,CAAC,GAAG,yBAAyB,CAAC,MAAM,CAAC,GAAG,aAAa,CAAC;YAC1F,CAAC;YAED,MAAM,UAAU,EAAE,CAAC,KAAK,EAAE,yBAAyB,CAAC,CAAC;YAErD,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;gBACpB,MAAM;YACV,CAAC;QACL,CAAC;QAED,EAAE,CAAC,OAAO,CAAC,kBAAkB,CAAC,CAAC;QAC/B,MAAM,UAAU,EAAE,EAAE,CAAA;QAEpB,OAAO,EAAE,CAAC;IACd,CAAC;IAGD;;;;;;;;;OASG;IACO,QAAQ,CACd,EAAa,EACb,EAAa,EACb,SAAgC,EAChC,aAA6B,EAC7B,WAAsD;QAKtD,IAAI,MAAiB,CAAC;QAEtB,+BAA+B;QAC/B,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,CAAC,aAAa,CAAC,GAAG,EAAE;YACjD,4DAA4D;YAC5D,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,EAAE,EAAE;gBACpB,QAAQ,EAAE,IAAI;gBACd,GAAG,WAAW;aACjB,CAAc,CAAC;YAEhB,6EAA6E;YAC7E,EAAE,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;YAEhB,MAAM,IAAI,GAAG,SAAS;gBAClB,CAAC,CAAC,aAAa,CAAC,EAAE,EAAE,MAAM,CAAC,CAAC,GAAG,CAAC,SAAS,CAAC;gBAC1C,CAAC,CAAC,aAAa,CAAC,EAAE,EAAE,MAAM,CAAC,CAAC;YAEhC,OAAO,IAAI,CAAC,IAAI,EAAe,CAAC;QACpC,CAAC,CAAC,CAAC;QAEH,kBAAkB;QAClB,IAAI,CAAC,SAAS,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC;QAErC,OAAO;YACH,MAAM,EAAE,MAAO;YACf,IAAI;SACP,CAAC;IACN,CAAC;IAGQ,OAAO,CAAC,IAAyB;QACtC,IAAI,IAAI,CAAC,IAAI,IAAI,yBAAyB,EAAE,CAAC;YACzC,MAAM,KAAK,CAAC,gIAAgI,CAAC,CAAA;QACjJ,CAAC;QAED,KAAK,CAAC,OAAO,CAAC,IAAI,CAAC,CAAC;IACxB,CAAC;IAGD;;;;;;;;;OASG;IACI,KAAK,CAAC,QAAQ,CAAC,KAAkB,EAAE,QAA0B,EAAE,SAA8C;QAChH,IAAI,QAAQ,CAAC,IAAI,IAAI,QAAQ,CAAC,iBAAiB,EAAE,CAAC;YAC9C,MAAM,KAAK,CAAC,sBAAsB,IAAI,CAAC,IAAI,mDAAmD,QAAQ,CAAC,iBAAiB,GAAG,CAAC,CAAC;QACjI,CAAC;QAED,IAAI,CAAC,cAAc,GAAG,KAAK,CAAC;QAE5B,IAAI,aAAa,GAAgB,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC,CAAC,CAAgB,CAAC,CAAC,kDAAkD;QAEtI,OAAO,CAAC,IAAI,CAAC,cAAc,IAAI,QAAQ,CAAC,IAAI,GAAG,QAAQ,CAAC,iBAAiB,EAAE,CAAC;YACxE,qEAAqE;YACrE,MAAM,UAAU,GAAG,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE,CAAC,IAAI,CAAC,gBAAgB,CAAC,aAAa,EAAE,QAAQ,CAAC,CAAC,CAAC;YAEjF,wDAAwD;YACxD,MAAM,oBAAoB,GAAG,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE,CAAC,UAAU,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YACpE,MAAM,SAAS,CAAC,oBAAoB,CAAC,CAAC;YAEtC,oBAAoB,CAAC,OAAO,EAAE,CAAC;YAE/B,aAAa,CAAC,OAAO,EAAE,CAAC;YACxB,aAAa,GAAG,UAAU,CAAC;QAC/B,CAAC;QAED,EAAE,CAAC,OAAO,CAAC,aAAa,CAAC,CAAC;IAC9B,CAAC;IAGD;;;;;;;OAOG;IACI,gBAAgB,CAAC,KAAkB,EAAE,QAA0B;QAClE,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,IAAI,CAAC,EAAE,CAAC;YACtB,MAAM,KAAK,CAAC,8BAA8B,IAAI,CAAC,IAAI,0CAA0C,CAAC,CAAC;QACnG,CAAC;QAED,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,qDAAqD;YACrD,MAAM,UAAU,GAAG,IAAI,CAAC,KAAK,CAAC,KAAK,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAc,CAAC;YAEzE,MAAM,CAAC,UAAU,EAAE,eAAe,EAAE,UAAU,CAAC,GAAG,UAAU,CAAC,KAAK,CAAC;YAEnE,qBAAqB;YACrB,MAAM,UAAU,GAAG,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,eAAe,GAAG,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,UAAU,EAAE,CAAC,EAAE,UAAU,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAA;YAEvG,OAAO,UAAyB,CAAC;QACrC,CAAC,CAAC,CAAA;IACN,CAAC;IAGD,IAAI,cAAc;QACd,OAAO,IAAI,CAAC,eAAe,CAAC;IAChC,CAAC;IAGD,IAAI,cAAc,CAAC,IAAa;QAC5B,IAAI,CAAC,eAAe,GAAG,IAAI,CAAC;IAChC,CAAC;;AAKL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,QAAQ,CAAC,CAAC"}

package/dist/src/models/u_net.d.ts

@@ -1,40 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
-export interface UNetArgs {
-    /**
-     * The starting number of filters.
-     */
-    filters: number;
-    /**
-     * The number of categories. For binary segmentation, `units=1`.
-     */
-    units: number;
-    /**
-     * The activation of the final output convolution layer. Defaults to `sigmoid` if `categories=1`, else `softmax`.
-     */
-    activation?: ActivationIdentifier;
-    /**
-     * The depth of the U-Net or the number of contractions and the number of expansions.
-     */
-    depth: number;
-    /**
-     * Adds residual connections to transform the model into a ResUNet. Defaults to `false`.
-     */
-    residual?: boolean;
-    /**
-     * Adds batch normalization to convolutions. Best used for batched inputs. Defaults to `false`.
-     */
-    batchNorm?: boolean;
-    /**
-     * Set the unbatched input shape of the U-Net in the format `[height, width, channels]`. Defaults to `[null, null, 3]`. If set, only channels is mandatory.
-     */
-    inputShape?: [number | null, number | null, number];
-}
-export type UNetModelArgs = UNetArgs & Omit<tf.SequentialArgs, "layers">;
-export declare class UNetModel extends tf.Sequential {
-    constructor(args: UNetModelArgs);
-    summary(lineLength?: number, positions?: number[], printFn?: (message?: any, ...optionalParams: any[]) => void): void;
-}
-export declare function createUNet({ filters, depth, units, activation, residual, batchNorm, inputShape }: UNetModelArgs): tf.LayersModel;
-export declare function loadUNetModel(pathOrIOHandler: string | tf.io.IOHandler, options?: tf.io.LoadOptions): Promise<tf.LayersModel>;
-//# sourceMappingURL=u_net.d.ts.map

package/dist/src/models/u_net.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"u_net.d.ts","sourceRoot":"","sources":["../../../src/models/u_net.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,oBAAoB,EAAE,MAAM,6DAA6D,CAAC;AAGxG,MAAM,WAAW,QAAQ;IACrB;;OAEG;IACH,OAAO,EAAE,MAAM,CAAC;IAChB;;OAEG;IACH,KAAK,EAAE,MAAM,CAAC;IACd;;OAEG;IACH,UAAU,CAAC,EAAE,oBAAoB,CAAC;IAClC;;OAEG;IACH,KAAK,EAAE,MAAM,CAAC;IACd;;OAEG;IACH,QAAQ,CAAC,EAAE,OAAO,CAAC;IACnB;;OAEG;IACH,SAAS,CAAC,EAAE,OAAO,CAAC;IACpB;;OAEG;IACH,UAAU,CAAC,EAAE,CAAC,MAAM,GAAG,IAAI,EAAE,MAAM,GAAG,IAAI,EAAE,MAAM,CAAC,CAAC;CACvD;AAGD,MAAM,MAAM,aAAa,GAAG,QAAQ,GAAG,IAAI,CAAC,EAAE,CAAC,cAAc,EAAE,QAAQ,CAAC,CAAC;AAGzE,qBAAa,SAAU,SAAQ,EAAE,CAAC,UAAU;gBAE5B,IAAI,EAAE,aAAa;IAsBtB,OAAO,CAAC,UAAU,CAAC,EAAE,MAAM,EAAE,SAAS,CAAC,EAAE,MAAM,EAAE,EAAE,OAAO,CAAC,EAAE,CAAC,OAAO,CAAC,EAAE,GAAG,EAAE,GAAG,cAAc,EAAE,GAAG,EAAE,KAAK,IAAI,GAAG,IAAI;CAIjI;AAGD,wBAAgB,UAAU,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,KAAK,EAAE,UAAU,EAAE,QAAgB,EAAE,SAAiB,EAAE,UAA4B,EAAE,EAAE,aAAa,kBA4CjJ;AAGD,wBAAsB,aAAa,CAAC,eAAe,EAAE,MAAM,GAAG,EAAE,CAAC,EAAE,CAAC,SAAS,EAAE,OAAO,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,WAAW,2BAOzG"}