From b6c18b986996e8856483aff30c731acd4d79696f Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Fri, 24 Apr 2026 11:08:40 +0200 Subject: Add logits script --- gpt2/scripts/logits.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 gpt2/scripts/logits.py 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) -- cgit v1.2.3