File size: 1,148 Bytes
cbe4d4c
 
ec89d20
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cbe4d4c
ec89d20
 
 
 
 
 
 
 
cbe4d4c
ec89d20
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import evaluate
from evaluate.utils import launch_gradio_widget
import evaluate
from evaluate.utils import launch_gradio_widget
import gradio as gr

# Define the list of available models
available_models = {
    "unitary/toxic-bert": "Bert Base Model",
    "facebook/roberta-hate-speech-dynabench-r4-target": "Facebook Model",
    "mskov/roberta-base-toxicity": "Roberta Finetuned Model",
    "mskov/distilbert-base-toxicity" : "Distillbert Finetuned Model"
}

# Create a Gradio interface with a radio button to select the model
def classify_toxicity(text, selected_model):
    # Load the selected model
    module = evaluate.load("toxicity", selected_model)
    results = module.compute(predictions=[text])
    toxicity_score = results["toxicity"][0]
    return f"Toxicity Score ({available_models[selected_model]}): {toxicity_score:.4f}"

iface = gr.Interface(
    fn=classify_toxicity,
    inputs=["text", gr.Radio(available_models, type="value", label="Select Model")],
    outputs="text",
    live=True,
    title="Toxicity Classifier",
    description="Select a model and enter text to classify its toxicity.",
)

launch_gradio_widget(iface)