engrossed elsewhere for the foreseeable, hence my absence. should the… — Bit x Wisdom — TG.ME

engrossed elsewhere for the foreseeable, hence my absence. should the day come, i'll resurface to scribe a line or two..

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_NAME = "cross-encoder/nli-deberta-v3-base"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME).to(DEVICE)
model.eval()

FREE_HYPOTHESIS = "The person is free and available."
BUSY_HYPOTHESIS = "The person is busy and unavailable."

ID_TO_LABEL = {
    int(index): label.lower()
    for index, label in model.config.id2label.items()
}

ENTAILMENT_ID = next(
    index for index, label in ID_TO_LABEL.items()
    if "entail" in label
)

CONTRADICTION_ID = next(
    index for index, label in ID_TO_LABEL.items()
    if "contrad" in label
)

ABSENCE_STATE = "BUSY"


@torch.inference_mode()
def _nli_scores(text, hypothesis):
    encoded = tokenizer(
        text,
        hypothesis,
        return_tensors="pt",
        truncation=True,
        max_length=128
    )

    encoded = {
        key: value.to(DEVICE)
        for key, value in encoded.items()
    }

    probabilities = torch.softmax(
        model(**encoded).logits,
        dim=-1
    )[0]

    return (
        probabilities[ENTAILMENT_ID].item(),
        probabilities[CONTRADICTION_ID].item()
    )


def status(text):
    if not isinstance(text, str) or not text.strip():
        return ABSENCE_STATE

    free_entailment, free_contradiction = _nli_scores(
        text,
        FREE_HYPOTHESIS
    )

    busy_entailment, busy_contradiction = _nli_scores(
        text,
        BUSY_HYPOTHESIS
    )

    free_score = free_entailment + busy_contradiction
    busy_score = busy_entailment + free_contradiction

    return "FREE" if free_score > busy_score else ABSENCE_STATE