import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification
)

from labels import LABELS

MODEL_PATH = "toxiguard-bert"

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH
)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_PATH
)

model.to(device)
model.eval()


def predict_toxicity(text):

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=256,
        padding=True
    )

    inputs = {
        key: value.to(device)
        for key, value in inputs.items()
    }

    with torch.no_grad():

        outputs = model(**inputs)

    probs = torch.sigmoid(
        outputs.logits
    ).cpu().numpy()[0]

    results = {}

    for i, prob in enumerate(probs):

        results[LABELS[i]] = round(
            float(prob), 4
        )

    return results