onnx-converter: new npm workspace to convert GPT2 from ONNX to TFJS

Author: JulienVig

cli/src/hellaswag_gpt.ts

@@ -1,50 +1,104 @@
+// import fs from 'fs';
+import fsPromise from 'node:fs/promises';
+
+import { dirname } from 'path';
+import { fileURLToPath } from 'url';
+import { parse } from 'ts-command-line-args'
+
 import '@tensorflow/tfjs-node';
 import fs from 'node:fs';
 import path from 'node:path';
-import { Tokenizer, models } from '@epfml/discojs';
+import { models, serialization, Tokenizer } from '@epfml/discojs';
 import { loadHellaSwag } from '@epfml/discojs-node';
+// import { AutoTokenizer } from '@xenova/transformers';
 
-const logFile = path.join('..', 'datasets', 'LogFile_hellaswag.txt');
-const logLines: string[] = [];
+const __dirname = dirname(fileURLToPath(import.meta.url));
 
+const logLines: string[] = [];
 function log(message: string) {
     console.log(message);
     logLines.push(message);
 }
 
-const hellaswagDataset: models.HellaSwagDataset = await loadHellaSwag(-1)
-
-async function evaluateTFJS(tokenizer: Tokenizer) {
-    const model = new models.GPT({ seed: 42 });
-    log('Evaluating TFJS GPT on HellaSwag...');
+async function evaluateModel(model: models.GPT | models.ONNXModel, numDataPoints = -1) {
+    const hellaswagDataset: models.HellaSwagDataset = await loadHellaSwag(numDataPoints)
+    const tokenizer = await Tokenizer.from_pretrained('Xenova/gpt2');
+    log('Starting the HellaSwag benchmark...');
 
     const start = Date.now();
-    const accuracy = await models.evaluate_hellaswag(model, tokenizer, hellaswagDataset, false);
+    const accuracy = await models.evaluate_hellaswag(model, tokenizer, hellaswagDataset, true);
     const duration = ((Date.now() - start) / 1000).toFixed(2);
 
-    log(`TFJS GPT Accuracy: ${(accuracy * 100).toFixed(2)}%`);
-    log(`TFJS GPT Evaluation Time: ${duration} seconds`);
+    log(`Final accuracy: ${(accuracy * 100).toFixed(2)}%`);
+    log(`Evaluation Time: ${duration} seconds`);
 }
 
-async function evaluateXenova(tokenizer: Tokenizer) {
-    const model = await models.ONNXModel.init_pretrained('Xenova/gpt2');
-    log('Evaluating Xenova GPT-2 (ONNX) on HellaSwag...');
+const ModelTypes = ['onnx', 'gpt-tfjs-random', 'gpt-tfjs-pretrained'] as const;
+type ModelType = typeof ModelTypes[number];
 
-    const start = Date.now();
-    const accuracy = await models.evaluate_hellaswag(model, tokenizer, hellaswagDataset, false);
-    const duration = ((Date.now() - start) / 1000).toFixed(2);
-
-    log(`Xenova GPT-2 Accuracy: ${(accuracy * 100).toFixed(2)}%`);
-    log(`Xenova GPT-2 Evaluation Time: ${duration} seconds`);
+interface HellaSwagArgs {
+    model: ModelType
+    numDataPoints: number
+    logFile: string
+    pretrainedModelPath: string
+    help?: boolean
 }
 
 async function main(): Promise<void> {
-    fs.writeFileSync(logFile, '', 'utf-8'); // Clear old log file
+    const defaultPretrainedModelPath = path.join(__dirname, "..", "..", "onnx-converter", "assets", "model.json")
+    const args = parse<HellaSwagArgs>({
+        model: {
+            type: (raw: string) => raw as ModelType,
+            description: `Model type, one of ${ModelTypes}`,
+            defaultValue: 'onnx'
+        },
+        numDataPoints: {
+            type: Number,
+            description: 'Number of HellaSwag datapoints to evaluate, set -1 for the whole benchmark',
+            defaultValue: -1
+        },
+        logFile: {
+            type: String,
+            description: 'Relative path to the log file, default to ./hellaswag.log', defaultValue: 'hellaswag.log'
+        },
+        pretrainedModelPath: {
+            type: String,
+            description: 'If specifying gpt-tfjs-pretrained, provide the relative path to the TF.js pretrained model',
+            defaultValue: defaultPretrainedModelPath
+        },
+        help: {
+            type: Boolean,
+            optional: true,
+            alias: 'h',
+            description: 'Prints this usage guide'
+        }
+    }, { helpArg: 'help' })
 
-    const tokenizer = await Tokenizer.from_pretrained('Xenova/gpt2');
-    await evaluateTFJS(tokenizer);
-    log('\n---\n');
-    await evaluateXenova(tokenizer);
+    const logFile = path.join(__dirname, args.logFile);
+    fs.writeFileSync(logFile, '', 'utf-8'); // Clear the log file
+
+    let model: | models.GPT | models.ONNXModel | undefined;
+    switch (args.model) {
+        case 'onnx':
+            log("Using ONNX pretrained model Xenova/gpt2")
+            model = await models.ONNXModel.init_pretrained('Xenova/gpt2');
+            break;
+            case 'gpt-tfjs-random':
+            log("Using GPT-TFJS with random initialization")
+            model = new models.GPT({ seed: 42 });
+            break;
+            case 'gpt-tfjs-pretrained':
+            log("Using GPT-TFJS with pretrained weights")
+            if (args.pretrainedModelPath === undefined) {
+                throw new Error("If choosing gpt-tfjs-pretrained, provide the relative path to the TF.js pretrained model `pretrainedModelPath")
+            }
+            const encodedModel = await fsPromise.readFile(args.pretrainedModelPath);
+            model = await serialization.model.decode(encodedModel) as models.GPT;
+            break;
+        default:
+            throw new Error(`Unrecognized model type: ${model}`);
+    } 
+    await evaluateModel(model, args.numDataPoints);
 
     fs.writeFileSync(logFile, logLines.join('\n'), 'utf-8');
     console.log(`\nResults written to ${logFile}`);

datasets/.gitignore

@@ -20,3 +20,6 @@
 
 # GDHF demo
 /tinder_dog/
+
+# HellaSwag benchmark
+hellaswag*

discojs/src/models/gpt/layers.spec.ts

@@ -174,8 +174,9 @@ describe('GPT Layers', () => {
       name: 'testCSA',
       contextLength: 5,
       nHead: 2,
-      nEmbd: 8,          // divisible by nHead, so head size = 4
-      dropout: 0.0,      // no dropout for deterministic tests
+      nEmbd: 8,      // divisible by nHead, so head size = 4
+      attnDrop: 0.0, // no dropout for deterministic tests
+      residDrop: 0.0,
       nLayer: 2,
       seed: 42
     };

discojs/src/models/hellaswag.ts

@@ -126,7 +126,7 @@ type ModelType = GPT | ONNXModel;
 export async function evaluate(
   model: ModelType,
   tokenizer: Tokenizer,
-  dataset: HellaSwagExample[],
+  dataset: HellaSwagDataset,
   print = true
 ): Promise<number> {
   let correct = 0;

onnx-converter/.gitignore

@@ -0,0 +1,3 @@
+node_modules
+assets
+dist

onnx-converter/README.md

@@ -0,0 +1,7 @@
+Structure:
+1. Read the ONNX model from Xenova's repository
+2. Use onnx.js protobuf to read the file and iterate through the layers: https://github.com/microsoft/onnxruntime/blob/main/js/web/lib/onnxjs/
+3. Create a map from layer to weight and convert each weight to TF.js tensor
+4. Init a TF.js model with the loaded weights and export the model
+
+Run `npm run convert_onnx` to create GPT-tfjs `model.json` file in the `./assets/` folder.

onnx-converter/package.json

@@ -0,0 +1,28 @@
+{
+  "name": "onnx-converter",
+  "private": true,
+  "type": "module",
+  "main": "dist/gpt2_from_onnx.js",
+  "scripts": {
+    "convert_onnx": "npm run build && node dist/convert_onnx.js",
+    "build": "tsc && cp -r src/protobuf dist/protobuf",
+    "lint": "npx eslint .",
+    "test": ": nothing"
+  },
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "@epfml/discojs-node": "*",
+    "@eslint/compat": "^1.4.0",
+    "@eslint/eslintrc": "^3.3.1",
+    "@eslint/js": "^9.36.0",
+    "globals": "^16.4.0",
+    "onnxruntime-web": "^1.22.0",
+    "server": "*",
+    "tslib": "2"
+  },
+  "devDependencies": {
+    "nodemon": "3",
+    "ts-command-line-args": "2"
+  }
+}