@epfml/discojs 3.0.1-p20250402145848.0 → 3.0.1-p20250625140656.0

package/dist/models/gpt/layers.d.ts

@@ -1,5 +1,72 @@
 import * as tf from '@tensorflow/tfjs';
 import type { GPTConfig } from './config.js';
+import type { ModelSize } from './config.js';
+/**
+ * Defines a range, from 0 to T, that is used to create positional embeddings
+ */
+export declare class Range extends tf.layers.Layer {
+    static readonly className = "Range";
+    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
+    call(input: tf.Tensor | tf.Tensor[], kwargs: Record<string, unknown>): tf.Tensor | tf.Tensor[];
+}
+export type CausalSelfAttentionConfig = ConstructorParameters<typeof tf.layers.Layer>[0] & Record<'contextLength' | 'nHead' | 'nEmbd' | 'dropout' | 'nLayer' | 'seed', number>;
+export declare class CausalSelfAttention extends tf.layers.Layer {
+    private readonly config;
+    static readonly className = "CausalSelfAttention";
+    private readonly nHead;
+    private readonly nEmbd;
+    private readonly nLayer;
+    private readonly dropout;
+    private readonly seed;
+    private readonly mask;
+    cAttnKernel?: tf.LayerVariable;
+    cAttnBias?: tf.LayerVariable;
+    cProjKernel?: tf.LayerVariable;
+    cProjBias?: tf.LayerVariable;
+    constructor(config: CausalSelfAttentionConfig);
+    build(): void;
+    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
+    getConfig(): tf.serialization.ConfigDict;
+    call(input: tf.Tensor | tf.Tensor[], kwargs: Record<string, unknown>): tf.Tensor;
+    dense(x: tf.Tensor, kernel: tf.LayerVariable, bias: tf.LayerVariable): tf.Tensor;
+    splitHeads(x: tf.Tensor, B: number, T: number, nHead: number): tf.Tensor;
+    applyCausalMask(att: tf.Tensor, T: number): tf.Tensor;
+    computeAttention(q: tf.Tensor, k: tf.Tensor, training: boolean, T: number): tf.Tensor;
+}
+/**
+ * GELU with tanh approximate
+ * GELU(x) = x * 0.5 * (1 + Tanh[sqrt(2/π) * (x + 0.044715 * x^3)])
+ *
+ * https://pytorch.org/docs/stable/generated/torch.nn.GELU.html
+ */
+export declare class GELU extends tf.layers.Layer {
+    static readonly className = "GELU";
+    constructor();
+    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
+    call(input: tf.Tensor | tf.Tensor[], kwargs: Record<string, unknown>): tf.Tensor | tf.Tensor[];
+}
+export type MLPConfig = ConstructorParameters<typeof tf.layers.Layer>[0] & Required<ModelSize> & Record<'contextLength' | 'residDrop' | 'nLayer' | 'seed', number>;
+export declare function MLP(config: MLPConfig): tf.LayersModel;
+/**
+ * LanguageModelEmbedding is a layer that combines the token embeddings and the language modeling head
+ * I.e. LMEmbedding is used to translate token indices into token embeddings
+ * as well as to project embeddings back into token indices
+ * The GPT2 model uses the same embedding matrix for both the token embeddings and the language modeling head
+ * Because Tensorflow.js doesn't offer an easy weight sharing mechanism, we need to define a custom layer
+ * that can be used for both the token embeddings and the language modeling head.
+ * In the GPT2 model definition, this layers corresponds to wte and lm_head (which reuses wte)
+ */
+export declare class LMEmbedding extends tf.layers.Layer {
+    private readonly vocabSize;
+    private readonly nEmbd;
+    private readonly seed;
+    static readonly className = "LMEmbedding";
+    embeddings?: tf.LayerVariable;
+    constructor(vocabSize: number, nEmbd: number, seed: number);
+    build(): void;
+    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
+    call(input: tf.Tensor | tf.Tensor[], kwargs: Record<string, unknown>): tf.Tensor | tf.Tensor[];
+}
 /**
  * The GPTArchitecture specifically defines a GPT forward pass, i.e.,
  * what are the inputs, the successive transformer blocks and the outputs. It is then

package/dist/models/gpt/layers.js

@@ -4,7 +4,7 @@ const debug = createDebug("discojs:models:gpt:layers");
 /**
  * Defines a range, from 0 to T, that is used to create positional embeddings
  */
-class Range extends tf.layers.Layer {
+export class Range extends tf.layers.Layer {
     static className = 'Range';
     computeOutputShape(inputShape) {
         return inputShape;
@@ -56,7 +56,7 @@ class LogLayer extends tf.layers.Layer {
     }
 }
 tf.serialization.registerClass(LogLayer);
-class CausalSelfAttention extends tf.layers.Layer {
+export class CausalSelfAttention extends tf.layers.Layer {
     config;
     static className = 'CausalSelfAttention';
     nHead;
@@ -86,8 +86,7 @@ class CausalSelfAttention extends tf.layers.Layer {
     }
     build() {
         // key, query, value projections for all heads, but in a batch
-        this.cAttnKernel = this.addWeight('c_attn.weight', [this.nEmbd, 3 * this.nEmbd], 'float32', tf.initializers.randomNormal({ mean: 0, stddev: 0.02, seed: this.seed }) // use same init as GPT2
-        );
+        this.cAttnKernel = this.addWeight('c_attn.weight', [this.nEmbd, 3 * this.nEmbd], 'float32', tf.initializers.randomNormal({ mean: 0, stddev: 0.02, seed: this.seed }));
         this.cAttnBias = this.addWeight('c_attn.bias', [3 * this.nEmbd], 'float32', tf.initializers.zeros());
         // output projection
         this.cProjKernel = this.addWeight('c_proj.kernel', [this.nEmbd, this.nEmbd], 'float32', 
@@ -97,7 +96,9 @@ class CausalSelfAttention extends tf.layers.Layer {
         // https://github.com/karpathy/build-nanogpt/blob/6104ab1b53920f6e2159749676073ff7d815c1fa/train_gpt2.py#L103
         // https://youtu.be/l8pRSuU81PU?si=5GcKfi_kPgLgvtg2&t=4640
         tf.initializers.randomNormal({
-            mean: 0, stddev: 0.02 * Math.sqrt(2 * this.nLayer), seed: this.seed
+            mean: 0,
+            stddev: 0.02 * Math.sqrt(2 * this.nLayer),
+            seed: this.seed
         }));
         this.cProjBias = this.addWeight('c_proj.bias', [this.nEmbd], 'float32', tf.initializers.zeros());
     }
@@ -122,59 +123,72 @@ class CausalSelfAttention extends tf.layers.Layer {
                 input = input[0];
             }
             this.invokeCallHook(input, kwargs);
-            const dense = (x, kernel, bias) => {
-                // TODO: use broadcasting when tfjs will support backpropagating through broadcasting
-                const k = kernel.read().expandDims(0).tile([x.shape[0], 1, 1]);
-                const m = x.matMul(k);
-                return tf.add(m, bias.read());
-            };
+            // --- Use helper methods below to build the computation ---     
             // Apply attention weights to inputs as one big matrix which is then split into the
             // query, key and value submatrices
             // nHead is "number of heads", hs is "head size", and C (number of channels) = n_embd = nHead * hs
-            // e.g. in GPT-2 (124M), nHead = 12, hs = 64, so nHead * hs = C = 768 channels in the Transformer
-            const cAttn = dense(input, this.cAttnKernel, this.cAttnBias);
+            // e.g. in GPT-2 (124M), nHead = 12, hs = 64, so nHead * hs = C = 768 channels in the Transformer      const cAttn = dense(input, this.cAttnKernel, this.cAttnBias);
+            const cAttn = this.dense(input, this.cAttnKernel, this.cAttnBias);
             let [q, k, v] = tf.split(cAttn, 3, -1);
             // Follow naming conventions in https://github.com/karpathy/build-nanogpt/
             const [B, T, C] = k.shape; // batch size, sequence length, embedding dimensionality (number of channels)
-            const splitHeads = (x) => tf.transpose(tf.reshape(x, [B, T, this.nHead, C / this.nHead]), // (B, T, nHead, head size)
-            [0, 2, 1, 3] // (B, nHead, T, hs)
-            );
-            q = splitHeads(q); // (B, nHead, T, hs)
-            k = splitHeads(k); // (B, nHead, T, hs)
-            v = splitHeads(v); // (B, nHead, T, hs)
+            // Split into attention heads.
+            q = this.splitHeads(q, B, T, this.nHead);
+            k = this.splitHeads(k, B, T, this.nHead);
+            v = this.splitHeads(v, B, T, this.nHead);
             // Scaled self attention: query @ key / sqrt(hs)
             // Matrix representing the token-to-token attention (B, nHead, T, T)
-            let att = tf.mul(tf.matMul(q, k, false, true), // (B, nHead, T, hs) x (B, nHead, hs, T) -> (B, nHead, T, T)
-            tf.div(1, tf.sqrt(tf.cast(k.shape[k.shape.length - 1], 'float32'))) // 1 / sqrt(hs)
-            );
-            /**
-             * The next operations apply attention only on the past tokens, which is
-             * essentially a weighted average of the past tokens with complicated weights,
-             * it relies on a mask to not "pay any attention" to future tokens
-             */
-            // mask is lower triangular matrix filled with 1
-            const mask = this.mask.slice([0, 0], [T, T]); // (T, T)
-            // 1 - mask                   => upper triangular matrix filled with 1
-            // (1 - mask) * -10^9         => upper triangular matrix filled with -inf
-            // att + ((1 - mask) * -10^9) => lower triangular part is the same as the `att` matrix
-            //                               upper triangular part is -inf
-            att = tf.add(att, tf.mul(tf.sub(1, mask), -1e9)); // (B, nHead, T, T)
-            // applying softmax zeroes out the upper triangular part (softmax(-inf) = 0)
-            // i.e., zeroes out future tokens's attention weights
-            // and creates a probability distribution for the lower triangular
-            // (attention weights of past tokens). The probability distribution ensures
-            // that the attention weights of past tokens for a particular token sum to one
-            att = tf.softmax(att, -1);
-            att = kwargs.training === true ? tf.dropout(att, this.dropout, undefined, this.seed) : att;
+            const att = this.computeAttention(q, k, kwargs.training === true, T);
             // This is where the (attention-)weighted sum of past values is performed
             let y = tf.matMul(att, v); // (B, nHead, T, T) x (B, nHead, T, hs) -> (B, nHead, T, hs)
             y = tf.transpose(y, [0, 2, 1, 3]); // (B, T, nHead, hs)
             y = tf.reshape(y, [B, T, C]); // (B, T, C = nHead * hs)
-            y = dense(y, this.cProjKernel, this.cProjBias); // output projection (B, T, C)
+            y = this.dense(y, this.cProjKernel, this.cProjBias); // output projection (B, T, C)
             y = kwargs.training === true ? tf.dropout(y, this.dropout, undefined, this.seed) : y;
             return y;
         });
     }
+    // --- Helper Methods ---
+    dense(x, kernel, bias) {
+        const k = kernel.read().expandDims(0).tile([x.shape[0], 1, 1]);
+        const m = x.matMul(k);
+        return tf.add(m, bias.read());
+    }
+    splitHeads(x, B, T, nHead) {
+        return tf.transpose(tf.reshape(x, [B, T, nHead, (x.shape[2] ?? 0) / nHead]), [0, 2, 1, 3]);
+    }
+    applyCausalMask(att, T) {
+        // mask is lower triangular matrix filled with 1
+        const mask = this.mask.slice([0, 0], [T, T]);
+        // 1 - mask                   => upper triangular matrix filled with 1
+        // (1 - mask) * -10^9         => upper triangular matrix filled with -inf
+        // att + ((1 - mask) * -10^9) => lower triangular part is the same as the `att` matrix
+        //                               upper triangular part is -inf
+        return tf.add(att, tf.mul(tf.sub(1, mask), -1e9)); // (B, nHead, T, T)
+    }
+    computeAttention(q, k, training, T) {
+        /**
+        * The next operations apply attention only on the past tokens, which is
+        * essentially a weighted average of the past tokens with complicated weights,
+        * it relies on a mask to not "pay any attention" to future tokens
+       */
+        const headSize = k.shape[k.shape.length - 1];
+        // Scaled self attention: query @ key / sqrt(hs)
+        // Matrix representing the token-to-token attention (B, nHead, T, T)
+        let att = tf.matMul(q, k, false, true); // (B, nHead, T, hs) x (B, nHead, hs, T) -> (B, nHead, T, T)
+        att = tf.mul(att, tf.div(1, tf.sqrt(tf.cast(headSize, 'float32')))); // 1 / sqrt(hs)
+        att = this.applyCausalMask(att, T);
+        // applying softmax zeroes out the upper triangular part (softmax(-inf) = 0)
+        // i.e., zeroes out future tokens's attention weights
+        // and creates a probability distribution for the lower triangular
+        // (attention weights of past tokens). The probability distribution ensures
+        // that the attention weights of past tokens for a particular token sum to one
+        att = tf.softmax(att, -1);
+        if (training) {
+            att = tf.dropout(att, this.dropout, undefined, this.seed);
+        }
+        return att;
+    }
 }
 tf.serialization.registerClass(CausalSelfAttention);
 /**
@@ -183,7 +197,7 @@ tf.serialization.registerClass(CausalSelfAttention);
  *
  * https://pytorch.org/docs/stable/generated/torch.nn.GELU.html
  */
-class GELU extends tf.layers.Layer {
+export class GELU extends tf.layers.Layer {
     static className = 'GELU';
     constructor() {
         super({});
@@ -210,7 +224,7 @@ class GELU extends tf.layers.Layer {
     }
 }
 tf.serialization.registerClass(GELU);
-function MLP(config) {
+export function MLP(config) {
     return tf.sequential({ layers: [
             tf.layers.dense({
                 name: config.name + `.mlp.c_fc`,
@@ -298,7 +312,7 @@ function TransformerBlock(conf) {
  * that can be used for both the token embeddings and the language modeling head.
  * In the GPT2 model definition, this layers corresponds to wte and lm_head (which reuses wte)
  */
-class LMEmbedding extends tf.layers.Layer {
+export class LMEmbedding extends tf.layers.Layer {
     vocabSize;
     nEmbd;
     seed;

package/dist/models/hellaswag.d.ts

@@ -0,0 +1,31 @@
+import { GPT } from './index.js';
+import { PreTrainedTokenizer } from '@xenova/transformers';
+import { ONNXModel } from './onnx.js';
+export declare const HELLASWAG_URL = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl";
+/**
+ * Represents a single example from the HellaSwag dataset.
+ *
+ * ctx - The context sentence or paragraph that sets up the situation.
+ * endings - An array of four possible continuations of the context.
+ * label - The index (0–3) of the correct ending in the `endings` array.
+ */
+export interface HellaSwagExample {
+    ctx: string;
+    endings: string[];
+    label: number;
+}
+export type HellaSwagDataset = HellaSwagExample[];
+type Tokenizer = PreTrainedTokenizer;
+type ModelType = GPT | ONNXModel;
+/**
+ * Evaluates the model on a given HellaSwag dataset.
+ *
+ * @param model - The model to evaluate (GPT or ONNXModel)
+ * @param tokenizer - The tokenizer to use
+ * @param dataset - An array of HellaSwagExample to evaluate on
+ * @param limit - Number of examples to evaluate (default: all)
+ * @param print - Whether to print results (default: true)
+ * @returns The accuracy of the model on the dataset
+ */
+export declare function evaluate(model: ModelType, tokenizer: Tokenizer, dataset: HellaSwagExample[], print?: boolean): Promise<number>;
+export {};

package/dist/models/hellaswag.js

@@ -0,0 +1,120 @@
+import * as tf from '@tensorflow/tfjs';
+import { GPT } from './index.js';
+import { tokenize } from '../processing/text.js';
+import { List } from 'immutable';
+export const HELLASWAG_URL = 'https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl';
+// Computes the log likelihood of the input sequence using the tfjs model
+// The input sequence is expected to be a concatenation of the context and the ending
+// The function computes the log likelihood of each ending and returns the one with the loss of each ending
+// Sources:
+// https://github.com/karpathy/build-nanogpt/blob/master/hellaswag.py
+//https://www.youtube.com/watch?v=l8pRSuU81PU
+async function computeLogLikelihood(gpt, inputIds, ctxLength) {
+    const lossTensor = tf.tidy(() => {
+        // Convert input sequence to shape [1, seq_len]
+        const inputTensor = tf.tensor2d([inputIds], [1, inputIds.length], 'int32');
+        // Get model logits: [1, seq_len, vocab_size]
+        const logits3D = gpt.extract().predict(inputTensor);
+        // Shift logits to align with next-token targets
+        const shiftedLogits = logits3D.slice([0, 0, 0], [1, inputIds.length - 1, -1]);
+        // Target tokens (next tokens), same length as shifted logits
+        const shiftedTargets = inputIds.slice(1);
+        const targetTensor = tf.tensor1d(shiftedTargets, 'int32');
+        // One-hot encode targets for cross-entropy loss
+        const oneHotLabels = tf.oneHot(targetTensor, shiftedLogits.shape[2]);
+        // Compute per-token cross-entropy log-probabilities (unnormalized loss)
+        const logProbs = tf.losses.softmaxCrossEntropy(oneHotLabels, shiftedLogits.squeeze());
+        // Create a mask to only include loss after the context length
+        const mask = tf.tensor1d(inputIds.map((_, i) => (i >= ctxLength ? 1 : 0)), 'float32').slice(1);
+        // Apply the mask and average over the selected tokens
+        const masked = logProbs.mul(mask);
+        const loss = masked.sum().div(mask.sum());
+        return loss;
+    });
+    const lossNumber = await lossTensor.array();
+    if (typeof lossNumber !== 'number') {
+        throw new Error('got multiple loss');
+    }
+    return lossNumber;
+}
+// Computes the log likelihood of the input sequence using the ONNX model
+// The input sequence is expected to be a concatenation of the context and the ending
+// The function computes the log likelihood of each ending and returns the one with the loss of each ending
+// Sources:
+// https://github.com/karpathy/build-nanogpt/blob/master/hellaswag.py
+// https://www.youtube.com/watch?v=l8pRSuU81PU 
+async function computeONNXLogLikelihood(model, inputIds, ctxLength) {
+    const batchInput = List([List(inputIds)]); // [1, seq_len]
+    // Run model to get logits: flattened [T * V]
+    const logitsTensor = await model.getLogits(batchInput);
+    const logits = logitsTensor.data;
+    const [_B, T, V] = logitsTensor.dims;
+    // Reshape flattened logits into [T][V]
+    const reshaped = Array.from({ length: T }, (_, t) => logits.slice(t * V, (t + 1) * V));
+    // Shift targets (next-token prediction)
+    const targets = inputIds.slice(1); // length = T - 1
+    const logitsShifted = reshaped.slice(0, T - 1); // also length = T - 1
+    // Compute per-token cross-entropy loss manually
+    const losses = logitsShifted.map((logit, i) => {
+        const maxLogit = Math.max(...logit); // for numerical stability
+        const exp = logit.map(x => Math.exp(x - maxLogit));
+        const sumExp = exp.reduce((a, b) => a + b, 0);
+        const probs = exp.map(e => e / sumExp); // softmax
+        return -Math.log(probs[targets[i]]); // cross-entropy loss
+    });
+    // Create a binary mask for non-context tokens
+    const mask = inputIds.map((_, i) => (i >= ctxLength ? 1 : 0)).slice(1);
+    // Apply the mask to the losses
+    const maskedLosses = losses.map((l, i) => l * mask[i]);
+    // Average the masked losses
+    const totalLoss = maskedLosses.reduce((a, b) => a + b, 0);
+    const sum = mask.reduce((a, b) => a + b, 0);
+    return totalLoss / (sum || 1); // avoid division by 0
+}
+/**
+ * Evaluates the model on a given HellaSwag dataset.
+ *
+ * @param model - The model to evaluate (GPT or ONNXModel)
+ * @param tokenizer - The tokenizer to use
+ * @param dataset - An array of HellaSwagExample to evaluate on
+ * @param limit - Number of examples to evaluate (default: all)
+ * @param print - Whether to print results (default: true)
+ * @returns The accuracy of the model on the dataset
+ */
+export async function evaluate(model, tokenizer, dataset, print = true) {
+    let correct = 0;
+    let total = 0;
+    for (const example of dataset) {
+        const endingTokens = example.endings.map(e => tokenize(tokenizer, example.ctx + ' ' + e, {
+            truncation: true,
+            max_length: 128
+        }).toArray());
+        const ctxTokens = tokenize(tokenizer, example.ctx, {
+            truncation: true,
+            max_length: 128
+        }).toArray();
+        let losses = [];
+        if (model instanceof GPT) {
+            losses = await Promise.all(endingTokens.map(e => computeLogLikelihood(model, e, ctxTokens.length)));
+        }
+        else {
+            losses = await Promise.all(endingTokens.map(e => computeONNXLogLikelihood(model, e, ctxTokens.length)));
+        }
+        const pred = losses.indexOf(Math.min(...losses));
+        if (pred === example.label)
+            correct++;
+        total++;
+        if (print) {
+            console.log(`\nExample #${total}`);
+            console.log(`Context: ${example.ctx}`);
+            example.endings.forEach((end, i) => {
+                console.log(`  ${i}: ${end}  (loss: ${losses[i].toFixed(4)})${i === example.label ? ' <-- correct' : ''}${i === pred ? ' <-- picked' : ''}`);
+            });
+            const accuracy_temp = correct / total;
+            console.log(`\n Accuracy on ${total} examples: ${(accuracy_temp * 100).toFixed(2)}%`);
+        }
+    }
+    const accuracy = correct / total;
+    console.log(`\nFinal accuracy on ${total} examples: ${(accuracy * 100).toFixed(2)}%`);
+    return accuracy;
+}

package/dist/models/index.d.ts

@@ -1,6 +1,9 @@
 export { Model } from './model.js';
 export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js";
 export { GPT } from './gpt/index.js';
+export { ONNXModel } from './onnx.js';
 export { GPTConfig } from './gpt/config.js';
+export { evaluate as evaluate_hellaswag } from './hellaswag.js';
 export { TFJS } from './tfjs.js';
 export { getTaskTokenizer } from './tokenizer.js';
+export { evaluate, HellaSwagDataset, HellaSwagExample, HELLASWAG_URL } from './hellaswag.js';

package/dist/models/index.js

@@ -1,5 +1,8 @@
 export { Model } from './model.js';
 export { EpochLogs } from "./logs.js";
 export { GPT } from './gpt/index.js';
+export { ONNXModel } from './onnx.js';
+export { evaluate as evaluate_hellaswag } from './hellaswag.js';
 export { TFJS } from './tfjs.js';
 export { getTaskTokenizer } from './tokenizer.js';
+export { evaluate, HELLASWAG_URL } from './hellaswag.js';

package/dist/models/onnx.d.ts

@@ -0,0 +1,19 @@
+import { Tensor } from '@xenova/transformers';
+import { Model } from './index.js';
+import type { WeightsContainer } from '../index.js';
+import { List } from 'immutable';
+import type { GenerationConfig as TFJSGenerationConfig } from './gpt/config.js';
+import type { Batched, DataFormat } from "../index.js";
+export declare class ONNXModel extends Model<'text'> {
+    #private;
+    private model;
+    private constructor();
+    static init_pretrained(modelName?: string): Promise<ONNXModel>;
+    getConfig(): Record<string, unknown>;
+    predict(batch: Batched<DataFormat.ModelEncoded["text"][0]>, options?: Partial<TFJSGenerationConfig>): Promise<Batched<DataFormat.ModelEncoded["text"][1]>>;
+    getLogits(batch: List<List<number>>): Promise<Tensor>;
+    train(): AsyncGenerator<never, never>;
+    get weights(): WeightsContainer;
+    set weights(_: WeightsContainer);
+    [Symbol.dispose](): void;
+}

package/dist/models/onnx.js

@@ -0,0 +1,71 @@
+import { AutoModelForCausalLM, Tensor } from '@xenova/transformers';
+import { Model } from './index.js';
+import { List } from 'immutable';
+import { DefaultGenerationConfig } from './gpt/config.js';
+export class ONNXModel extends Model {
+    model;
+    constructor(model) {
+        super();
+        this.model = model;
+    }
+    static async init_pretrained(modelName = 'Xenova/gpt2') {
+        const model = await AutoModelForCausalLM.from_pretrained(modelName);
+        return new ONNXModel(model);
+    }
+    getConfig() {
+        return this.model.config;
+    }
+    async predict(batch, options) {
+        const config = Object.assign({}, DefaultGenerationConfig, options);
+        return List(await Promise.all(batch.map(tokens => this.#predictSingle(tokens, config))));
+    }
+    async #predictSingle(tokens, config) {
+        const contextLength = this.model.config.max_position_embeddings ?? 1024;
+        const truncated = tokens.slice(-contextLength).toArray();
+        if (truncated.length === 0) {
+            throw new Error('Token list is empty. Cannot run generate().');
+        }
+        const input_ids = new Tensor('int64', truncated.map(BigInt), [1, truncated.length]);
+        const output = await this.model.generate(input_ids, {
+            max_new_tokens: 1,
+            temperature: config.temperature,
+            do_sample: config.doSample,
+            top_k: config.topk,
+        });
+        if (!Array.isArray(output) || output.length === 0 || !Array.isArray(output[0])) {
+            throw new Error('ONNX model.generate() did not return valid sequences.');
+        }
+        const predicted_id = output[0].at(-1);
+        return Number(predicted_id);
+    }
+    async getLogits(batch) {
+        const input_ids_array = batch.toArray().map(seq => seq.toArray());
+        const attention_mask_array = input_ids_array.map((seq) => new Array(seq.length).fill(1));
+        const input_ids_flat = input_ids_array.flat();
+        const attention_mask_flat = attention_mask_array.flat();
+        const shape = [input_ids_array.length, input_ids_array[0].length];
+        // use BigInt for int64 compatibility
+        const input_ids = new Tensor('int64', input_ids_flat.map(BigInt), shape);
+        const attention_mask = new Tensor('int64', attention_mask_flat.map(BigInt), shape);
+        // run model forward
+        const outputs = await this.model.forward({ input_ids, attention_mask });
+        return outputs.logits;
+    }
+    async *train() {
+        await Promise.resolve(); // dummy await
+        const yieldFlag = false;
+        if (yieldFlag)
+            yield undefined; // satisfy 'require-yield'
+        throw new Error('Training not supported for ONNX models');
+    }
+    get weights() {
+        throw new Error('Weights access not supported in ONNX models');
+    }
+    set weights(_) {
+        throw new Error('Weights setting not supported in ONNX models');
+    }
+    [Symbol.dispose]() {
+        // Dispose of the model to free up memory
+        void this.model.dispose();
+    }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epfml/discojs",
-  "version": "3.0.1-p20250402145848.0",
+  "version": "3.0.1-p20250625140656.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",