@epfml/discojs 3.0.1-p20250429140233.0 → 3.0.1-p20250701123931.0

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-p20250429140233.0",
+  "version": "3.0.1-p20250701123931.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",