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() | |