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.