Spaces:
Sleeping
Sleeping
# Import gradio and transformers libraries | |
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
# Load the small deberta models for hate and offensive speech detection | |
hate_model = AutoModelForSequenceClassification.from_pretrained("KoalaAI/HateSpeechDetector") | |
hate_tokenizer = AutoTokenizer.from_pretrained("KoalaAI/HateSpeechDetector") | |
offensive_model = AutoModelForSequenceClassification.from_pretrained("KoalaAI/OffensiveSpeechDetector") | |
offensive_tokenizer = AutoTokenizer.from_pretrained("KoalaAI/OffensiveSpeechDetector") | |
# Define a function that takes an input text and returns the scores from the models | |
def get_scores(text): | |
# Tokenize and encode the input text | |
hate_input = hate_tokenizer(text, return_tensors="pt") | |
offensive_input = offensive_tokenizer(text, return_tensors="pt") | |
# Get the logits from the models | |
hate_logits = hate_model(**hate_input).logits | |
offensive_logits = offensive_model(**offensive_input).logits | |
# Apply softmax to get probabilities | |
hate_probs = hate_logits.softmax(dim=1) | |
offensive_probs = offensive_logits.softmax(dim=1) | |
# Get the labels from the models | |
hate_labels = hate_model.config.id2label | |
offensive_labels = offensive_model.config.id2label | |
# Format the output as a dictionary of scores | |
output = {} | |
output["Hate speech"] = {hate_labels[i]: round(p.item(), 4) for i, p in enumerate(hate_probs[0])} | |
output["Offensive speech"] = {offensive_labels[i]: round(p.item(), 4) for i, p in enumerate(offensive_probs[0])} | |
return output | |
# Create a gradio interface with a text input and a json output | |
iface = gr.Interface(fn=get_scores, inputs="text", outputs="json") | |
# Launch the interface | |
iface.launch() | |