@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

package/{src/losses/dice.ts → dist/losses/dice.js}

@@ -1,156 +1,93 @@
 import * as tf from "@tensorflow/tfjs";
 import { categoricalCrossentropy, binaryCrossentropy } from "@tensorflow/tfjs-layers/dist/losses";
-
 const epsilon = 1e-7;
-
 const REDUCE_HW = [1, 2]; // reduce over width and height
 const REDUCE_BHW = [0, 1, 2]; // reduce over batch, width, height
 const REDUCE_BHWC = [0, 1, 2, 3]; // reduce all dimensions
-
-
 // Standard (Sorensen) Dice Loss
-export function diceBinaryStandard(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
-
+export function diceBinaryStandard(y_true, y_pred) {
     const y_true_flat = tf.reshape(y_true, [y_true.shape[0], -1]);
     const y_pred_flat = tf.reshape(y_pred, [y_pred.shape[0], -1]);
-
     const intersection = tf.sum(tf.mul(y_true_flat, y_pred_flat), 1);
     const union = tf.add(tf.sum(y_true_flat, 1), tf.sum(y_pred_flat, 1));
-
-    const dice = tf.div(
-        intersection.mul(2).add(epsilon),
-        union.add(epsilon)
-    );
-
+    const dice = tf.div(intersection.mul(2).add(epsilon), union.add(epsilon));
     return tf.scalar(1).sub(dice);
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceBinaryStandard, "name", { value: "diceBinaryStandard", configurable: false });
-
-
 // https://github.com/keras-team/keras/blob/v3.3.3/keras/src/losses/losses.py#L1983-L2010
-export function diceBinaryGlobal(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+export function diceBinaryGlobal(y_true, y_pred) {
     const y_true_flat = tf.reshape(y_true, [-1]);
     const y_pred_flat = tf.reshape(y_pred, [-1]);
-
     const intersection = tf.sum(tf.mul(y_true_flat, y_pred_flat));
     const union = tf.add(tf.sum(y_true_flat), tf.sum(y_pred_flat));
-
-    const dice = tf.div(
-        intersection.mul(2).add(epsilon),
-        union.add(epsilon)
-    );
-
+    const dice = tf.div(intersection.mul(2).add(epsilon), union.add(epsilon));
     return tf.scalar(1).sub(dice);
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceBinaryGlobal, "name", { value: "diceBinaryGlobal", configurable: false });
-
-
-export function diceCategoricalStandard(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+export function diceCategoricalStandard(y_true, y_pred) {
     const intersection = tf.sum(tf.mul(y_true, y_pred), REDUCE_HW);
     const union = tf.add(y_true, y_pred).sum(REDUCE_HW);
-
-    const dice = tf.div(
-        intersection.mul(2).add(epsilon),
-        union.add(epsilon)
-    );
-
+    const dice = tf.div(intersection.mul(2).add(epsilon), union.add(epsilon));
     return tf.scalar(1).sub(tf.mean(dice, -1));
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceCategoricalStandard, "name", { value: "diceCategoricalStandard", configurable: false });
-
-
-export function diceCategoricalGeneralized(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
-
+export function diceCategoricalGeneralized(y_true, y_pred) {
     // this is done twice so we calculate it once
     const y_true_sum = y_true.sum(REDUCE_BHW);
-
     const weighting = tf.div(1, y_true_sum.square().add(epsilon));
-
     const intersection = tf.sum(tf.mul(y_true, y_pred), REDUCE_BHW).mul(weighting).sum();
     const union = tf.add(y_true_sum, y_pred.sum(REDUCE_BHW)).mul(weighting).sum();
-
-    const dice = tf.div(
-        intersection.mul(2).add(epsilon),
-        union.add(epsilon)
-    );
-
+    const dice = tf.div(intersection.mul(2).add(epsilon), union.add(epsilon));
     return tf.scalar(1).sub(dice);
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceCategoricalGeneralized, "name", { value: "diceCategoricalGeneralized", configurable: false });
-
-
-export function diceCategoricalGlobal(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
-
+export function diceCategoricalGlobal(y_true, y_pred) {
     const intersection = tf.sum(tf.mul(y_true, y_pred), REDUCE_BHWC);
     const union = tf.add(tf.sum(y_true, REDUCE_BHWC), tf.sum(y_pred, REDUCE_BHWC));
-
-    const dice = tf.div(
-        intersection.mul(2).add(epsilon),
-        union.add(epsilon)
-    );
-
+    const dice = tf.div(intersection.mul(2).add(epsilon), union.add(epsilon));
     return tf.scalar(1).sub(dice);
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceCategoricalGlobal, "name", { value: "diceCategoricalGlobal", configurable: false });
-
-
 /**
  * Calculates the Sorensen-Dice coefficient and the binary cross entropy losses.
  * Both have equal weight.
- * 
+ *
  * @param y_true the label tensor
  * @param y_pred the prediction tensor (not sparse)
  * @returns a tensor of shape `[ batch ]`
  */
-export function diceBinaryCrossentropy(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+export function diceBinaryCrossentropy(y_true, y_pred) {
     // reduce cross entropy shape from [B, H, W] to [B] to match dice
     const bce = binaryCrossentropy(y_true, y_pred).mean(REDUCE_HW);
     const dice = diceBinaryStandard(y_true, y_pred);
-
     return tf.add(bce.mul(0.5), dice.mul(0.5));
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceBinaryCrossentropy, "name", { value: "diceBinaryCrossentropy", configurable: false });
-
-
 /**
  * Calculates the Sorensen-Dice coefficient and the categorical cross entropy losses.
  * Both have equal weight. Expects dense (non-sparse) label tensors.
- * 
+ *
  * This does not support sparse tensors because TFJS's
  * sparseCategoricalCrossentropy loss onehots the label
  * and calls categoricalCrossentropy. See
  * https://github.com/tensorflow/tfjs/blob/0fc04d958ea592f3b8db79a8b3b497b5c8904097/tfjs-layers/src/losses.ts#L143-L146
- * 
- * @param y_true the label  
+ *
+ * @param y_true the label
  * @param y_pred the prediction tensor (not sparse)
  * @returns a tensor of shape `[ batch ]`
  */
-export function diceCategoricalCrossentropy(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+export function diceCategoricalCrossentropy(y_true, y_pred) {
     // reduce cross entropy shape from [B, H, W] to [B] to match dice
     const cce = categoricalCrossentropy(y_true, y_pred).mean(REDUCE_HW);
     const dice = diceCategoricalStandard(y_true, y_pred);
-
     return tf.add(cce.mul(0.5), dice.mul(0.5));
 }
-
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(diceCategoricalCrossentropy, "name", { value: "diceCategoricalCrossentropy", configurable: false });
+//# sourceMappingURL=dice.js.map

package/dist/losses/dice.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"dice.js","sourceRoot":"","sources":["../../src/losses/dice.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,uBAAuB,EAAE,kBAAkB,EAAE,MAAM,qCAAqC,CAAC;AAElG,MAAM,OAAO,GAAG,IAAI,CAAC;AAErB,MAAM,SAAS,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,+BAA+B;AACzD,MAAM,UAAU,GAAG,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,mCAAmC;AACjE,MAAM,WAAW,GAAG,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,wBAAwB;AAG1D,gCAAgC;AAChC,MAAM,UAAU,kBAAkB,CAAC,MAAiB,EAAE,MAAiB;IAEnE,MAAM,WAAW,GAAG,EAAE,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAC9D,MAAM,WAAW,GAAG,EAAE,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAE9D,MAAM,YAAY,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,WAAW,EAAE,WAAW,CAAC,EAAE,CAAC,CAAC,CAAC;IACjE,MAAM,KAAK,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,WAAW,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,GAAG,CAAC,WAAW,EAAE,CAAC,CAAC,CAAC,CAAC;IAErE,MAAM,IAAI,GAAG,EAAE,CAAC,GAAG,CACf,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,CAAC,EAChC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,CACrB,CAAC;IAEF,OAAO,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;AAClC,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,kBAAkB,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,oBAAoB,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAGxG,yFAAyF;AACzF,MAAM,UAAU,gBAAgB,CAAC,MAAiB,EAAE,MAAiB;IACjE,MAAM,WAAW,GAAG,EAAE,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;IAC7C,MAAM,WAAW,GAAG,EAAE,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;IAE7C,MAAM,YAAY,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,WAAW,EAAE,WAAW,CAAC,CAAC,CAAC;IAC9D,MAAM,KAAK,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,WAAW,CAAC,EAAE,EAAE,CAAC,GAAG,CAAC,WAAW,CAAC,CAAC,CAAC;IAE/D,MAAM,IAAI,GAAG,EAAE,CAAC,GAAG,CACf,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,CAAC,EAChC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,CACrB,CAAC;IAEF,OAAO,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;AAClC,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,gBAAgB,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,kBAAkB,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAGpG,MAAM,UAAU,uBAAuB,CAAC,MAAiB,EAAE,MAAiB;IACxE,MAAM,YAAY,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,EAAE,SAAS,CAAC,CAAC;IAC/D,MAAM,KAAK,GAAG,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;IAEpD,MAAM,IAAI,GAAG,EAAE,CAAC,GAAG,CACf,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,CAAC,EAChC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,CACrB,CAAC;IAEF,OAAO,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;AAC/C,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,uBAAuB,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,yBAAyB,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAGlH,MAAM,UAAU,0BAA0B,CAAC,MAAiB,EAAE,MAAiB;IAE3E,6CAA6C;IAC7C,MAAM,UAAU,GAAG,MAAM,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC;IAE1C,MAAM,SAAS,GAAG,EAAE,CAAC,GAAG,CAAC,CAAC,EAAE,UAAU,CAAC,MAAM,EAAE,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC,CAAC;IAE9D,MAAM,YAAY,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,EAAE,UAAU,CAAC,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC,GAAG,EAAE,CAAC;IACrF,MAAM,KAAK,GAAG,EAAE,CAAC,GAAG,CAAC,UAAU,EAAE,MAAM,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC,GAAG,EAAE,CAAC;IAE9E,MAAM,IAAI,GAAG,EAAE,CAAC,GAAG,CACf,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,CAAC,EAChC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,CACrB,CAAC;IAEF,OAAO,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;AAClC,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,0BAA0B,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,4BAA4B,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAGxH,MAAM,UAAU,qBAAqB,CAAC,MAAiB,EAAE,MAAiB;IAEtE,MAAM,YAAY,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,EAAE,WAAW,CAAC,CAAC;IACjE,MAAM,KAAK,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,WAAW,CAAC,EAAE,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,WAAW,CAAC,CAAC,CAAC;IAE/E,MAAM,IAAI,GAAG,EAAE,CAAC,GAAG,CACf,YAAY,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,CAAC,EAChC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,CACrB,CAAC;IAEF,OAAO,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;AAClC,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,qBAAqB,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,uBAAuB,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAG9G;;;;;;;GAOG;AACH,MAAM,UAAU,sBAAsB,CAAC,MAAiB,EAAE,MAAiB;IACvE,iEAAiE;IACjE,MAAM,GAAG,GAAG,kBAAkB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;IAC/D,MAAM,IAAI,GAAG,kBAAkB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAEhD,OAAO,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;AAC/C,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,sBAAsB,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,wBAAwB,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAGhH;;;;;;;;;;;;GAYG;AACH,MAAM,UAAU,2BAA2B,CAAC,MAAiB,EAAE,MAAiB;IAC5E,iEAAiE;IACjE,MAAM,GAAG,GAAG,uBAAuB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;IACpE,MAAM,IAAI,GAAG,uBAAuB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAErD,OAAO,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;AAC/C,CAAC;AAGD,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,2BAA2B,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,6BAA6B,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC"}

package/dist/losses/index.d.ts

@@ -0,0 +1,2 @@
+export * from "./dice";
+//# sourceMappingURL=index.d.ts.map

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

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

package/dist/losses/index.js

@@ -0,0 +1,2 @@
+export * from "./dice";
+//# sourceMappingURL=index.js.map

package/dist/losses/index.js.map

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

package/dist/masks.d.ts

@@ -0,0 +1,20 @@
+import * as tf from "@tensorflow/tfjs";
+/**
+ * Generate a causal mask used in self-attention to prevent tokens from looking
+ * ahead. The values in the upper right portion of the mask matrix are set to
+ * -1e7 so that they have no impact during scaled dot product attention.
+ */
+export declare function causal(query_seq_length: number, key_seq_length: number): tf.Tensor<tf.Rank>;
+/**
+ * Generate a self-attention mask that prevents packed sequences from cross document
+ * boundaries and attending to each other. The result is a tensor of diagonally
+ * positioned blocks of zeroes (allow attention) and -1e7 values (prevent attention).
+ * The latter is scored zero during the scaled dot product attention's softmax operation.
+ *
+ * @param boundaries an array of ones (denotes start of a new sample or docment) and zeroes
+ *
+ * Example boundary of 3 samples that are packed into one:
+ * `[1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0]`
+ */
+export declare function packing(boundaries: Int32Array): tf.Tensor<tf.Rank>;
+//# sourceMappingURL=masks.d.ts.map

package/dist/masks.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"masks.d.ts","sourceRoot":"","sources":["../src/masks.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAGvC;;;;GAIG;AACH,wBAAgB,MAAM,CAAC,gBAAgB,EAAE,MAAM,EAAE,cAAc,EAAE,MAAM,sBAItE;AAGD;;;;;;;;;;GAUG;AACH,wBAAgB,OAAO,CAAC,UAAU,EAAE,UAAU,sBAc7C"}

package/{src/packing_mask.ts → dist/masks.js}

@@ -1,28 +1,37 @@
 import * as tf from "@tensorflow/tfjs";
-
+/**
+ * Generate a causal mask used in self-attention to prevent tokens from looking
+ * ahead. The values in the upper right portion of the mask matrix are set to
+ * -1e7 so that they have no impact during scaled dot product attention.
+ */
+export function causal(query_seq_length, key_seq_length) {
+    return tf.linalg.bandPart(tf.ones([query_seq_length, key_seq_length]), -1, 0)
+        .sub(1)
+        .mul(1e7);
+}
 /**
  * Generate a self-attention mask that prevents packed sequences from cross document
  * boundaries and attending to each other. The result is a tensor of diagonally
  * positioned blocks of zeroes (allow attention) and -1e7 values (prevent attention).
  * The latter is scored zero during the scaled dot product attention's softmax operation.
- * 
+ *
  * @param boundaries an array of ones (denotes start of a new sample or docment) and zeroes
- * 
- * Example boundary of 3 samples are packed into one:
+ *
+ * Example boundary of 3 samples that are packed into one:
  * `[1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0]`
  */
-export function generatePackingSelfAttentionMask(boundaries: Int32Array) {
+export function packing(boundaries) {
     // see images at
     // https://reddit.com/r/LocalLLaMA/comments/197efaz/training_llama_mistral_and_mixtralmoe_faster_with/
     return tf.tidy(() => {
         // cumsum transforms the tensor such that each sequence in the pack gets its own id,
         const partitions = tf.tensor1d(boundaries).cumsum();
-
         return partitions.expandDims(1)
             .equal(partitions.expandDims(0))
             .sub(1)
             .mul(1e7)
             // introduce a head dimension so it can be broadcasted
             .expandDims(0);
-    })
+    });
 }
+//# sourceMappingURL=masks.js.map

package/dist/masks.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"masks.js","sourceRoot":"","sources":["../src/masks.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAGvC;;;;GAIG;AACH,MAAM,UAAU,MAAM,CAAC,gBAAwB,EAAE,cAAsB;IACnE,OAAO,EAAE,CAAC,MAAM,CAAC,QAAQ,CAAC,EAAE,CAAC,IAAI,CAAC,CAAC,gBAAgB,EAAE,cAAc,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC;SACxE,GAAG,CAAC,CAAC,CAAC;SACN,GAAG,CAAC,GAAG,CAAC,CAAC;AAClB,CAAC;AAGD;;;;;;;;;;GAUG;AACH,MAAM,UAAU,OAAO,CAAC,UAAsB;IAC1C,gBAAgB;IAChB,sGAAsG;IACtG,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;QAChB,oFAAoF;QACpF,MAAM,UAAU,GAAG,EAAE,CAAC,QAAQ,CAAC,UAAU,CAAC,CAAC,MAAM,EAAE,CAAC;QAEpD,OAAO,UAAU,CAAC,UAAU,CAAC,CAAC,CAAC;aAC1B,KAAK,CAAC,UAAU,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;aAC/B,GAAG,CAAC,CAAC,CAAC;aACN,GAAG,CAAC,GAAG,CAAC;YACT,sDAAsD;aACrD,UAAU,CAAC,CAAC,CAAC,CAAC;IACvB,CAAC,CAAC,CAAA;AACN,CAAC"}

package/dist/metrics.d.ts

@@ -0,0 +1,20 @@
+import { Tensor } from "@tensorflow/tfjs";
+/**
+ * Applies the recall metric with the prediction rounded based on a threshold
+ *
+ * @param y_true the label tensor
+ * @param y_pred the prediction tensor
+ * @param threshold threshold value to be considered a positive prediction, defaults to `0.5`
+ * @returns
+ */
+export declare function recall(y_true: Tensor, y_pred: Tensor, threshold?: number): Tensor<import("@tensorflow/tfjs-core").Rank>;
+/**
+ * Applies the precision metric with the prediction rounded based on a threshold
+ *
+ * @param y_true the label tensor
+ * @param y_pred the prediction tensor
+ * @param threshold threshold value to be considered a positive prediction, defaults to `0.5`
+ * @returns
+ */
+export declare function precision(y_true: Tensor, y_pred: Tensor, threshold?: number): Tensor<import("@tensorflow/tfjs-core").Rank>;
+//# sourceMappingURL=metrics.d.ts.map

package/dist/metrics.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"metrics.d.ts","sourceRoot":"","sources":["../src/metrics.ts"],"names":[],"mappings":"AAAA,OAAO,EAAW,MAAM,EAAE,MAAM,kBAAkB,CAAC;AAGnD;;;;;;;GAOG;AACH,wBAAgB,MAAM,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,SAAS,GAAE,MAAY,gDAE7E;AAKD;;;;;;;GAOG;AACH,wBAAgB,SAAS,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,SAAS,GAAE,MAAY,gDAEhF"}

package/{src/metrics.ts → dist/metrics.js}

@@ -1,32 +1,28 @@
-import { metrics, Tensor } from "@tensorflow/tfjs";
-
-
+import { metrics } from "@tensorflow/tfjs";
 /**
  * Applies the recall metric with the prediction rounded based on a threshold
- * 
+ *
  * @param y_true the label tensor
  * @param y_pred the prediction tensor
  * @param threshold threshold value to be considered a positive prediction, defaults to `0.5`
- * @returns 
+ * @returns
  */
-export function recall(y_true: Tensor, y_pred: Tensor, threshold: number = 0.5) {
+export function recall(y_true, y_pred, threshold = 0.5) {
     return metrics.recall(y_true, y_pred.greaterEqual(threshold));
 }
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(recall, "name", { value: "recall", configurable: false });
-
 /**
  * Applies the precision metric with the prediction rounded based on a threshold
- * 
+ *
  * @param y_true the label tensor
  * @param y_pred the prediction tensor
  * @param threshold threshold value to be considered a positive prediction, defaults to `0.5`
- * @returns 
+ * @returns
  */
-export function precision(y_true: Tensor, y_pred: Tensor, threshold: number = 0.5) {
+export function precision(y_true, y_pred, threshold = 0.5) {
     return metrics.precision(y_true, y_pred.greaterEqual(threshold));
 }
-
 // prevents minification of function name which TFJS relies on
 Object.defineProperty(precision, "name", { value: "precision", configurable: false });
+//# sourceMappingURL=metrics.js.map

package/dist/metrics.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"metrics.js","sourceRoot":"","sources":["../src/metrics.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,OAAO,EAAU,MAAM,kBAAkB,CAAC;AAGnD;;;;;;;GAOG;AACH,MAAM,UAAU,MAAM,CAAC,MAAc,EAAE,MAAc,EAAE,YAAoB,GAAG;IAC1E,OAAO,OAAO,CAAC,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,YAAY,CAAC,SAAS,CAAC,CAAC,CAAC;AAClE,CAAC;AAED,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,QAAQ,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC;AAEhF;;;;;;;GAOG;AACH,MAAM,UAAU,SAAS,CAAC,MAAc,EAAE,MAAc,EAAE,YAAoB,GAAG;IAC7E,OAAO,OAAO,CAAC,SAAS,CAAC,MAAM,EAAE,MAAM,CAAC,YAAY,CAAC,SAAS,CAAC,CAAC,CAAC;AACrE,CAAC;AAED,8DAA8D;AAC9D,MAAM,CAAC,cAAc,CAAC,SAAS,EAAE,MAAM,EAAE,EAAE,KAAK,EAAE,WAAW,EAAE,YAAY,EAAE,KAAK,EAAE,CAAC,CAAC"}

package/dist/models/gpt_model.d.ts

@@ -0,0 +1,94 @@
+import * as tf from "@tensorflow/tfjs";
+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;
+}
+/**
+ * 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 declare class GptModel extends LlmModel {
+    static className: string;
+    protected readonly numHeads: number;
+    protected readonly numLayers: number;
+    protected readonly embedDim: number;
+    protected readonly vocabSize: number;
+    protected readonly padToMultipleOf64: boolean;
+    protected readonly vocabSizePadded: number;
+    protected vocab_padding_mask?: tf.Tensor1D;
+    /**
+     * DO NOT add layers in the constructor or it will break tf.loadLayersModel().
+     * It should be done in build() instead.
+     */
+    constructor(args: GptModelArgs);
+    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;
+    };
+    /**
+     * 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: tf.Tensor2D, kv_cache: KvCacheContainer): tf.Tensor2D;
+    build(inputShape?: tf.Shape | tf.Shape[]): void;
+    dispose(): DisposeResult;
+    getConfig(): {
+        numHeads: number;
+        numLayers: number;
+        embedDim: number;
+        vocabSize: number;
+        vocabSizePadded: number;
+        padToMultipleOf64: boolean;
+    };
+}
+//# sourceMappingURL=gpt_model.d.ts.map

package/dist/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,eAAe,CAAC;AACpD,OAAO,EAAE,QAAQ,EAAE,KAAK,YAAY,EAAE,MAAM,qBAAqB,CAAC;AAClE,OAAO,EAAE,gBAAgB,EAAE,MAAM,aAAa,CAAC;AAC/C,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"}