kaptaan45/QaptaanLM-0.75B-ONNX

import os
import numpy as np
import onnxruntime as ort
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

model_id = "kaptaan45/QaptaanLM-0.75B-ONNX"

model_dir = snapshot_download(repo_id=model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

session = ort.InferenceSession(
    os.path.join(model_dir, "model.onnx"),
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

prompt = 'def is_palindrome(s: str) -> bool:\n    \"\"\"Return True if s is a palindrome.\"\"\"\n    '
input_ids = tokenizer.encode(prompt, return_tensors="np")

outputs = session.run(None, {"input_ids": input_ids})
logits = outputs[0]
next_token = int(np.argmax(logits[0, -1, :]))
print("Next Token ID:", next_token, "Decoded:", repr(tokenizer.decode([next_token])))