@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

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

@@ -1,89 +1,48 @@
 import * as tf from "@tensorflow/tfjs";
-import * as tfc from "@/.";
-import { type LossOrMetricFn } from "@/tfjs_types";
-import { LlmModel, type LlmModelArgs } from "@/models/llm_model";
-import { KvCacheContainer } from "@/kv_cache";
-import { type DisposeResult } from "@tensorflow/tfjs-layers/dist/engine/topology";
-
-
-export interface GptModelArgs extends LlmModelArgs {
-    /**
-     * Number of heads per attention layer.
-     */
-    numHeads: number;
-    /**
-     * Number of GPT decoder blocks.
-     */
-    numLayers: number;
-    /**
-     * The embedding size of each token.
-     */
-    embedDim: number;
-    /**
-     * The vocabulary size of the embedding layer and number of units of the output
-     * layer. This is also the tokenizer vocabulary size.
-     */
-    vocabSize: number;
-    /**
-     * Pad the embeddings' vocab size and output layer's units to the next nearest 
-     * multiple of 64 to optimize hardware efficiency. Defaults to `true`.
-     * 
-     * For example: if a tokenizer has 50,257 tokens, the model uses 50,304 for the
-     * vocab size and output units count.
-     */
-    padToMultipleOf64?: boolean;
-}
-
-
+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";
-
-    protected readonly numHeads: number;
-    protected readonly numLayers: number;
-    protected readonly embedDim: number;
-    protected readonly vocabSize: number;
-    protected readonly padToMultipleOf64: boolean;
-
+    numHeads;
+    numLayers;
+    embedDim;
+    vocabSize;
+    padToMultipleOf64;
     // this is kept for reproducibility and model history but is not important since
     // it can be calculated mathematically
-    protected readonly vocabSizePadded: number;
-
+    vocabSizePadded;
     // the amount to pad the embedding vocab size and dense output units count
-    protected vocab_padding_mask?: tf.Tensor1D;
-
-
+    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: GptModelArgs) {
+    constructor(args) {
         const { numHeads, numLayers, embedDim, vocabSize, padToMultipleOf64 = true, ...rest } = args;
-
         super({ name: "model", ...rest });
-
         this.numHeads = numHeads;
         this.numLayers = numLayers;
         this.embedDim = embedDim;
@@ -93,68 +52,47 @@ export class GptModel extends LlmModel {
             ? Math.ceil(this.vocabSize / 64) * 64
             : this.vocabSize;
     }
-
-
-    protected override fitBatch(
-        xs: tf.Tensor,
-        ys: tf.Tensor,
-        loss_mask: tf.Tensor | undefined,
-        loss_function: LossOrMetricFn,
-        other_masks?: { [key: string]: tf.Tensor | undefined }
-    ) {
-        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(() => {
             y_pred = this.apply(xs, {
                 training: true,
                 ...other_masks
-            }) as tf.Tensor;
-
+            });
             // 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() as tf.Scalar;
+            return loss.mean();
         });
-
         // backpropagation
         this.optimizer.applyGradients(grads);
-
         return {
-            y_pred: y_pred!,
+            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
      */
-    override predictNextToken(input: tf.Tensor2D, kv_cache: KvCacheContainer): tf.Tensor2D {
+    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 }) 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 = this.vocab_padding_mask != undefined
                 ? prediction
@@ -166,52 +104,40 @@ export class GptModel extends LlmModel {
                     .slice([0, sequence_length - 1, 0], [batch_size, 1, vocab_size])
                     .softmax()
                     .argMax(2);
-
-            return next_token as tf.Tensor2D;
-        })
+            return next_token;
+        });
     }
-
-
-    override build(inputShape?: tf.Shape | tf.Shape[]): void {
+    build(inputShape) {
         const actual_vocab_size = this.vocabSizePadded
             ? this.vocabSizePadded
             : this.padToMultipleOf64
                 ? Math.ceil(this.vocabSize / 64) * 64
-                : this.vocabSize
-
+                : 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(_ => tfc.gpt2DecoderBlock({ numHeads: this.numHeads, embedDim: this.embedDim })),
+                ...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))
+            ].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<tf.Tensor1D>(
-                // 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())
+            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);
     }
-
-
-    override dispose(): DisposeResult {
+    dispose() {
         this.vocab_padding_mask?.dispose();
         return super.dispose();
     }
-
-
-    override getConfig() {
+    getConfig() {
         const base_config = super.getConfig();
-
         const config = {
             numHeads: this.numHeads,
             numLayers: this.numLayers,
@@ -219,14 +145,10 @@ export class GptModel extends LlmModel {
             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/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,qBAAqB,CAAC;AAGlE,OAAO,EAAE,gBAAgB,EAAE,MAAM,6BAA6B,CAAC;AAgC/D;;;;;;;;;;;;;;;;;;;;;;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/models/index.d.ts

@@ -0,0 +1,7 @@
+import { LlmModel, type LlmModelArgs } from "./llm_model";
+import { GptModel, type GptModelArgs } from "../models/gpt_model";
+import { type UNetArgs } from "../models/u_net";
+export declare function llmModel(args: LlmModelArgs): LlmModel;
+export declare function gptModel(args: GptModelArgs): GptModel;
+export declare function unetModel(args: UNetArgs): import("@tensorflow/tfjs-layers").LayersModel;
+//# sourceMappingURL=index.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/models/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,QAAQ,EAAE,KAAK,YAAY,EAAE,MAAM,aAAa,CAAC;AAC1D,OAAO,EAAE,QAAQ,EAAE,KAAK,YAAY,EAAE,MAAM,qBAAqB,CAAC;AAClE,OAAO,EAAc,KAAK,QAAQ,EAAE,MAAM,iBAAiB,CAAC;AAG5D,wBAAgB,QAAQ,CAAC,IAAI,EAAE,YAAY,YAE1C;AAGD,wBAAgB,QAAQ,CAAC,IAAI,EAAE,YAAY,YAE1C;AAGD,wBAAgB,SAAS,CAAC,IAAI,EAAE,QAAQ,iDAEvC"}

package/dist/models/index.js

@@ -0,0 +1,13 @@
+import { LlmModel } from "./llm_model";
+import { GptModel } from "../models/gpt_model";
+import { createUNet } from "../models/u_net";
+export function llmModel(args) {
+    return new LlmModel(args);
+}
+export function gptModel(args) {
+    return new GptModel(args);
+}
+export function unetModel(args) {
+    return createUNet(args);
+}
+//# sourceMappingURL=index.js.map

package/dist/models/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/models/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,QAAQ,EAAqB,MAAM,aAAa,CAAC;AAC1D,OAAO,EAAE,QAAQ,EAAqB,MAAM,qBAAqB,CAAC;AAClE,OAAO,EAAE,UAAU,EAAiB,MAAM,iBAAiB,CAAC;AAG5D,MAAM,UAAU,QAAQ,CAAC,IAAkB;IACvC,OAAO,IAAI,QAAQ,CAAC,IAAI,CAAC,CAAC;AAC9B,CAAC;AAGD,MAAM,UAAU,QAAQ,CAAC,IAAkB;IACvC,OAAO,IAAI,QAAQ,CAAC,IAAI,CAAC,CAAC;AAC9B,CAAC;AAGD,MAAM,UAAU,SAAS,CAAC,IAAc;IACpC,OAAO,UAAU,CAAC,IAAI,CAAC,CAAC;AAC5B,CAAC"}

package/dist/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/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,eAAe,CAAC;AAE7D,OAAO,EAAE,gBAAgB,EAAE,MAAM,aAAa,CAAC;AAK/C,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"}