from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
import torch

# Hate Speech
tokenizer = AutoTokenizer.from_pretrained(
    "mrm8488/distilroberta-finetuned-tweets-hate-speech")
model = AutoModelForSequenceClassification.from_pretrained(
    "mrm8488/distilroberta-finetuned-tweets-hate-speech")


def classify_hatespeech(sentence):
    preprocessed_text = sentence.strip().replace("\n", "")
    inputs = tokenizer(preprocessed_text, return_tensors="pt")
    labels = torch.tensor([1]).unsqueeze(0)
    outputs = model(**inputs, labels=labels)
    logits = outputs.logits
    probs = torch.softmax(logits, dim=1)
    nice = torch.flatten(probs).detach().numpy()[0]
    return "{:.2f}".format(nice)