#!/usr/bin/python
# -*- coding: utf-8 -*-

import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel, AutoConfig
from transformers.onnx import FeaturesManager, convert, export
from pathlib import Path
import os

model_id = "./distilgpt2-base-pretrained-he"
export_folder = "tmp/onnx/"
file_name = "model.onnx"

print('Loading tokenizer...')
tokenizer = AutoTokenizer.from_pretrained(model_id)
print('Saving tokenizer to ', export_folder)
tokenizer.save_pretrained(export_folder)
print('Loading model...')
model = AutoModelForCausalLM.from_pretrained(model_id)

feature= "causal-lm"
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=feature)
onnx_config = model_onnx_config(model.config)

print("model_kind = {0}\nonx_config = {1}\n".format(model_kind, onnx_config))

onnx_path = Path(export_folder+file_name)

print('Exporting model to ', onnx_path)
onnx_inputs, onnx_outputs = export(tokenizer, model, onnx_config, onnx_config.default_onnx_opset, onnx_path)
print('Done')