import torch

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
)

MODEL_PATH = "."

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_PATH
)


def predict(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
    )

    with torch.no_grad():
        outputs = model(**inputs)

    probabilities = torch.softmax(
        outputs.logits,
        dim=-1
    )[0]

    predicted_id = torch.argmax(probabilities).item()

    predicted_label = model.config.id2label[predicted_id]
    confidence = probabilities[predicted_id].item()

    return predicted_label, confidence


while True:
    text = input("\nDescribe an accessibility issue (or 'quit'): ")

    if text.lower() == "quit":
        break

    label, confidence = predict(text)

    print(f"\nPrediction: {label}")
    print(f"Confidence: {confidence:.1%}")