blob: 4f413b2d4a8150641cdf1f38b8daef1fae3641ab (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
|
#!/usr/bin/env python3
import json
from optimum.onnxruntime import ORTModelForCausalLM
from transformers import AutoTokenizer
cpu_provider = "CPUExecutionProvider"
cuda_provider = "CUDAExecutionProvider"
model = ORTModelForCausalLM.from_pretrained(
"../models/base",
provider=cuda_provider,
local_files_only=True,
use_cache=False,
)
tokenizer = AutoTokenizer.from_pretrained("../models/base")
text = "The quick brown fox jumps over the lazy dog"
inputs = tokenizer(text, return_tensors="pt")
print(inputs)
outputs = model(**inputs)
logits = outputs.logits
export = logits.squeeze(0).detach().cpu().numpy().astype("float32")
export.tofile("logits.f32")
seq_len, vocab = export.shape
with open("shape.json", "w") as f:
json.dump({"seq_len": int(seq_len), "vocab": int(vocab)}, f)
|
