初始化项目,由ModelHub XC社区提供模型
Model: Norod78/distilgpt2-base-pretrained-he Source: Original Platform
This commit is contained in:
31
converters/convert2onnx.py
Normal file
31
converters/convert2onnx.py
Normal file
@@ -0,0 +1,31 @@
|
||||
#!/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')
|
||||
Reference in New Issue
Block a user