@stellarapp/tfjs-stellar 1.0.0 → 1.0.2

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

@@ -1,24 +1,8 @@
 import * as tf from "@tensorflow/tfjs";
-import * as tfc from "@/index";
 import { sparseCategoricalCrossentropy } from "@tensorflow/tfjs-layers/dist/losses";
-import { Dataset, type LossOrMetricFn } from "@/tfjs_types";
-import { generateCausalAttentionMask } from "@/utils";
-import { KvCacheContainer } from "@/kv_cache";
-
-
-// eslint-disable-next-line
-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;
-}
-
-
+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
@@ -26,42 +10,33 @@ interface DatasetArgs extends tf.TensorContainerObject {
  */
 export class LlmModel extends tf.Sequential {
     static className = "LlmModel";
-
-    private stopPredicting_: boolean = true;
-
-    constructor(args: LlmModelArgs) {
+    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"
      */
-    protected getMetricFunctions() {
+    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.
      */
-    protected getLossFunction(): LossOrMetricFn {
+    getLossFunction() {
         let loss = this.loss;
-
         if (Array.isArray(loss)) {
             loss = loss[0];
         }
-
         if (typeof loss == "string") {
             if (loss == "sparseCategoricalCrossentropy") {
                 return sparseCategoricalCrossentropy;
@@ -70,249 +45,174 @@ export class LlmModel extends tf.Sequential {
                     " Use categoricalCrossentropy instead. See" +
                     " https://github.com/tensorflow/tfjs/blob/0fc04d958ea592f3b8db79a8b3b497b5c8904097/tfjs-layers/src/losses.ts#L143-L146"); */
             }
-
-            const loss_id = loss as string;
-
-            const loss_fn =
-                ((tfc.losses as Record<string, any>)[loss_id] ??
-                    (tf.losses as Record<string, any>)[loss_id] ??
-                    (tf.metrics as Record<string, any>)[loss_id]) as LossOrMetricFn
-
+            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 {
+                return loss_fn;
+            }
+            else {
                 throw Error(`LlmModel.getLossFunction: ${loss_id} is not a valid loss function`);
             }
-        } else if (typeof 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
      */
-    override async fitDataset<T = DatasetArgs>(tfdataset: Dataset<T>, args: tf.ModelFitDatasetArgs<T>): Promise<any> {
+    async fitDataset(tfdataset, args) {
         this.stopTraining = false;
-
-        const dataset = tfdataset as tf.data.Dataset<DatasetArgs>;
+        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 as tf.CustomCallbackArgs ?? {};
-
+        const { onBatchBegin, onBatchEnd, onEpochBegin, onEpochEnd, onTrainBegin, onTrainEnd, } = callbacks ?? {};
         await onTrainBegin?.();
-
-        let cached_causal_mask: tf.Tensor | undefined = undefined;
-
+        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: { [metric: string]: number } = {};
-
+            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: { [metric: string]: number } = { batch };
-
+                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 = generateCausalAttentionMask(seq_length, 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_sum = metric_fn(ys, y_pred).mean();
                         const metric_value = (metric_sum.dataSync())[0];
-
-                        batch_metrics[metric_label] = metric_value// / batch_size;
+                        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(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?.()
-
+        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
      */
-    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;
-    } {
-        let y_pred: tf.Tensor;
-
+    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
-            }) as tf.Tensor;
-
+            });
             // 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() as tf.Scalar;
+            return loss.mean();
         });
-
         // backpropagation
         this.optimizer.applyGradients(grads);
-
         return {
-            y_pred: y_pred!,
+            y_pred: y_pred,
             loss
         };
     }
-
-
-    override compile(args: tf.ModelCompileArgs): void {
+    compile(args) {
         if (args.loss == "categoricalCrossentropy") {
-            throw Error(`LlmModel.compile: use sparseCategoricalCrossentropy loss (along with onehot encoded labels) instead of 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
      */
-    public async generate(input: tf.Tensor1D, kv_cache: KvCacheContainer, onPredict: (token: tf.Tensor) => Promise<void>) {
+    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.Tensor2D = tf.tidy(() => input.expandDims(0)) as tf.Tensor2D; // it's 2D because of the required batch dimension
-
+        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
@@ -321,35 +221,25 @@ export class LlmModel extends tf.Sequential {
      * position of `sentence_length` is returned as the next predicted
      * token.
      */
-    public predictNextToken(input: tf.Tensor2D, kv_cache: KvCacheContainer) {
+    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 }) as tf.Tensor;
-
+            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 as tf.Tensor2D;
-        })
+            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: boolean) {
+    set stopPredicting(stop) {
         this.stopPredicting_ = stop;
     }
-
 }
-
-
 tf.serialization.registerClass(LlmModel);
+//# sourceMappingURL=llm_model.js.map

package/dist/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,UAAU,CAAC;AAExD,OAAO,KAAK,MAAM,MAAM,WAAW,CAAC;AAKnC,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/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/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"}