@stellarapp/tfjs-stellar 1.0.0 → 1.0.2

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

@@ -0,0 +1 @@
+{"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

@@ -0,0 +1,154 @@
+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

@@ -0,0 +1 @@
+{"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

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

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

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

package/{src/models/index.ts → dist/src/models/index.js}

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

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

@@ -0,0 +1 @@
+{"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

@@ -0,0 +1,87 @@
+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

@@ -0,0 +1 @@
+{"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

@@ -0,0 +1,245 @@
+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

@@ -0,0 +1 @@
+{"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

@@ -0,0 +1,40 @@
+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

@@ -0,0 +1 @@
+{"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"}