Add logits script

1 files changed, 35 insertions, 0 deletions

diff --git a/gpt2/scripts/logits.py b/gpt2/scripts/logits.py
new file mode 100644
index 0000000..4f413b2
--- /dev/null
+++ b/gpt2/scripts/logits.py
@@ -0,0 +1,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)