Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from detoxify import Detoxify | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| def run_model(model_choice, text): | |
| if model_choice == "original": | |
| results = Detoxify("original").predict(text) | |
| elif model_choice == "unbiased": | |
| results = Detoxify("unbiased").predict(text) | |
| elif model_choice == "multilingual": | |
| results = Detoxify("multilingual").predict(text) | |
| elif model_choice == "Toxic-BERT": | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| "citizenlab/distilbert-base-multilingual-cased-toxicity" | |
| ) | |
| model = AutoModelForSequenceClassification.from_pretrained( | |
| "citizenlab/distilbert-base-multilingual-cased-toxicity" | |
| ) | |
| # Perform inference with the Toxic-BERT model using tokenizer and model | |
| encoded_input = tokenizer( | |
| text, padding=True, truncation=True, max_length=512, return_tensors="pt" | |
| ) | |
| logits = model(**encoded_input).logits | |
| probabilities = torch.sigmoid(logits).detach().cpu().numpy().tolist()[0] | |
| results = {"toxic": probabilities[0], "non_toxic": probabilities[1]} | |
| # Convert the predicted_labels to your desired output format | |
| return results | |
| model_choices = ["original", "unbiased", "multilingual", "Toxic-BERT"] | |
| input_textbox = gr.inputs.Textbox(label="Input Text") | |
| model_choice_dropdown = gr.inputs.Dropdown(choices=model_choices, label="Model Choice") | |
| output_text = gr.outputs.Textbox(label="Output Tags") | |
| iface = gr.Interface( | |
| fn=run_model, | |
| inputs=[model_choice_dropdown, input_textbox], | |
| outputs=output_text, | |
| title="Toxicity Detection App", | |
| description="Choose a model and input text to detect toxicity.", | |
| ) | |
| iface.launch() | |