blob: e3773e24322e90560338624591aa409637364016 (
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("../models/base")
|