Harikrishna-Srinivasan's picture
Upload folder using huggingface_hub
bdb6fd8 verified
Raw
History Blame
1.08 kB
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LABELS = ["Not Hate", "Hate"]
MODEL_ID = "."
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to(DEVICE)
model.eval()
def predict(text: str):
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=512,
token_type_ids=False
)
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
with torch.inference_mode():
logits = model(**inputs).logits
probs = torch.softmax(logits, dim=-1)[0]
not_hate, hate = probs[0].item(), probs[1].item()
return {
"Not Hate": not_hate,
"Hate": hate
}
gr.Interface(
fn=predict,
inputs=gr.Textbox(lines=3, placeholder="Enter the text"),
outputs=gr.Label(num_top_classes=2),
title="RoBERTa Hate Speech Classifier"
).launch()