blob: cf0cbb893bed4256605fad2b7fd344fd216c92e9 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
|
# /// script
# dependencies = [
# "torch",
# "transformers",
# "optimum",
# "optimum[onnxruntime]",
# ]
# ///
from transformers import AutoTokenizer, AutoModelForCausalLM
from optimum.onnxruntime import ORTModelForCausalLM
model_id = "gpt2"
model = ORTModelForCausalLM.from_pretrained(model_id, export=True, use_cache=True)
model.save_pretrained("onnx-gpt2")
|
