ONNX to Tensorflow.js conversion of GPT-2

[christinakopi]

Roadmap:

• ensure bijection between onnx and tfjs weights
• add progress indicator to hellaswag benchmark
• move filesystem logic to discojs-node and cli
• automate downloading onnx model from hub
• convert python script to nodejs and web-compatible logic
• experiment with webgpu for hellaswag
• experiment with converting tfjs back to onnx

In order to be able to run the load_weights.spec.ts you need to change the format of the pretrained weights from .onnx to .jsonl by running the following python script:

Before running the script make sure you have downloaded the decoder_model.onnx from the following link: https://huggingface.co/Xenova/gpt2/tree/main/onnx and you have added it in the ONNX_MODEL_PATH

!pip install onnx
!pip install onnxruntime

import onnx
import numpy as np
from onnx import numpy_helper
import json
import sys

ONNX_MODEL_PATH = "<path/to/decoder_model.onnx>"
OUTPUT_FILENAME = "<where the output file (gpt2_weights.jsonl) will be saved>"

try:
    model = onnx.load(ONNX_MODEL_PATH)
except FileNotFoundError:
    print(f"ERROR: The model file was not found at '{ONNX_MODEL_PATH}'")
    sys.exit(1)

print(f"Extracting weights from {ONNX_MODEL_PATH}...")

with open(OUTPUT_FILENAME, "w") as f:
    for initializer in model.graph.initializer:
        name = initializer.name
        array = numpy_helper.to_array(initializer)

        weight_object = {
            "key": name,
            "value": array.tolist()
        }

        f.write(json.dumps(weight_object) + "\n")

print(f"Successfully saved weights to {OUTPUT_FILENAME}")

After generating this file, place it under the following path: disco/discojs/gpt2_weights.jsonl

The table below shows the accuracy along with the time of the evaluation on the whole HellaSwag for each model tested:

Model	Accuracy	Eval Time (s)
TFJS GPT (gpt-nano)	24.67%	1390.25
Xenova GPT-2 (ONNX)	29.03%	22767.03
Loaded TFJS GPT (from ONNX)	28.41%	12523.59

Note: The loaded_hellaswag.spec.ts reproduces this result for the loaded model. For the other 2 models this can be reproduced using the hellaswag_gpt.ts script in cli/

[martinjaggi]

amazing work, thanks a lot!
later i'd be very curious to check if also the conversion in the reverse direction (to onnx) works fine.

[tharvik]

haa, that's great work, well done in hacking around the onnx protobuf!
nothing blocking (except the tsc --build thing), only a few nitpicks and questions here and there

[tharvik]

commented

[tharvik]

commented

[tharvik]

we don't need a log system for the CLI, we can simply output to the console, no?

[tharvik]

that was confusing

Suggested change

	case 'gpt-tfjs-random':
	log("Using GPT-TFJS with random initialization")
	model = new models.GPT({ seed: 42 });
	break;
	case 'gpt-tfjs-pretrained':
	case 'gpt-tfjs-random':
	log("Using GPT-TFJS with random initialization")
	model = new models.GPT({ seed: 42 });
	break;
	case 'gpt-tfjs-pretrained':

[tharvik]

casting isn't nice, especially for user interfaces, better to implement so value checking with a switch

[tharvik]

that look to me as more relevant field than using rawData, why not using it first (and drop the warning)?

[tharvik]

Suggested change

	const finalWeights = gptLayersModel.weights.map((weight, _i) => {
	const finalWeights = gptLayersModel.weights.map((weight) => {

[tharvik]

can we more simply store it in the current directory and so letting choose the user where to have it?

[tharvik]

we don't have tests anyway

Suggested change

"exclude": ["**/*.spec.ts"]

[tharvik]

single use, we can inline it IMO