from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
import onnxruntime as ort
import numpy as np

tokenizer = AutoTokenizer.from_pretrained(
    "palaciodata/NorBERTo-base-NLI-Multitask"
)

model_path = hf_hub_download(
    repo_id="palaciodata/NorBERTo-base-NLI-Multitask",
    filename="model.onnx"
)

session = ort.InferenceSession(model_path)

premissa = "Uma criança risonha está segurando uma pistola de água e sendo espirrada com água"
hipotese = "Uma criança está segurando uma pistola de água"

inputs = tokenizer(
    premissa,
    hipotese,
    return_tensors="np",
    max_length=256,
    truncation=True,
    padding="max_length"
)

logits, similarity = session.run(
    None,
    {
        "input_ids": inputs["input_ids"].astype(np.int64),
        "attention_mask": inputs["attention_mask"].astype(np.int64),
    }
)

classe = np.argmax(logits, axis=1)[0]

print(
    f"Classe: {classe}, Similaridade: {float(similarity[0]):.3f}"
)