Spaces:
Sleeping
Sleeping
| import time | |
| import gradio as gr | |
| from transformers import pipeline | |
| model_options = [ | |
| "Mezzo-Content-Guard-Large", | |
| "Mezzo-Content-Guard-Base", | |
| "Mezzo-Content-Guard-Small", | |
| "Mezzo-Content-Guard-Large-Preview" | |
| ] | |
| cached_models = {} | |
| def run_model(model_name, input_text): | |
| if not input_text.strip(): | |
| return {}, "0ms" | |
| model = cached_models.get(model_name) | |
| if not model: | |
| repo_id = f"RyanStudio/{model_name}" | |
| model = pipeline("text-classification", repo_id) | |
| cached_models[model_name] = model | |
| _ = model("warmup") | |
| start = time.time() | |
| raw_output = model(input_text, top_k=None) # get all classes | |
| latency = f"{round((time.time() - start) * 1000, 2)} ms" | |
| if isinstance(raw_output, list) and isinstance(raw_output[0], list): | |
| raw_output = raw_output[0] | |
| elif isinstance(raw_output, dict): | |
| raw_output = [raw_output] | |
| scores = {item["label"]: item["score"] for item in raw_output} | |
| return scores, latency | |
| def run_comparison(model_name_1, model_name_2, input_text): | |
| if not input_text.strip(): | |
| return {}, {}, "0ms", "0ms" | |
| scores_1, latency_1 = run_model(model_name_1, input_text) | |
| scores_2, latency_2 = run_model(model_name_2, input_text) | |
| return scores_1, scores_2, latency_1, latency_2 | |
| custom_css = """ | |
| #container { max-width: 900px; margin: auto; padding-top: 20px; } | |
| .output-stats { font-weight: bold; color: #555; } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="container"): | |
| gr.Markdown("# 🛡️ Mezzo Guard – Content Moderation") | |
| gr.Markdown("Analyze text for 5 content categories (**Content Guard**).") | |
| comparison_toggle = gr.Checkbox(label="Compare 2 Models", value=False) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| text_input = gr.Textbox( | |
| label="Input Prompt", | |
| placeholder="Enter the text you want to screen...", | |
| lines=6, | |
| max_lines=15 | |
| ) | |
| with gr.Group(visible=True) as single_model_group: | |
| model_dropdown = gr.Dropdown( | |
| label="Model", | |
| choices=model_options, | |
| value=model_options[0], | |
| interactive=True | |
| ) | |
| with gr.Group(visible=False) as comparison_group: | |
| model_dropdown_1 = gr.Dropdown( | |
| label="Model 1", | |
| choices=model_options, | |
| value=model_options[0], | |
| interactive=True | |
| ) | |
| model_dropdown_2 = gr.Dropdown( | |
| label="Model 2", | |
| choices=model_options, | |
| value=model_options[1], | |
| interactive=True | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear", variant="secondary") | |
| run_button = gr.Button("Analyze Prompt", variant="primary") | |
| with gr.Column(scale=1): | |
| with gr.Group(visible=True) as single_output_group: | |
| label_output = gr.Label(label="Classification Result", num_top_classes=6) | |
| latency_output = gr.Textbox(label="Latency", interactive=False, elem_classes="output-stats") | |
| # Comparison output | |
| with gr.Group(visible=False) as comparison_output_group: | |
| with gr.Row(): | |
| with gr.Column(): | |
| label_output_1 = gr.Label(label="Model 1 Results", num_top_classes=6) | |
| latency_output_1 = gr.Textbox(label="Model 1 Latency", interactive=False, elem_classes="output-stats") | |
| with gr.Column(): | |
| label_output_2 = gr.Label(label="Model 2 Results", num_top_classes=6) | |
| latency_output_2 = gr.Textbox(label="Model 2 Latency", interactive=False, elem_classes="output-stats") | |
| gr.Markdown("### Performance Info") | |
| gr.HTML( | |
| "<small>Model weights are cached after the first run. Content Guard returns scores for violence, hate-speech, sexual, toxic and self-harm</small>" | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["Hello how are you?"], | |
| ["I am so horny right now"], | |
| ["I hate faggots"], | |
| ["I want to kill myself"], | |
| ["I fucking hate you cause ur a cunt"], | |
| ["I want to bomb a hospital right now"] | |
| ], | |
| inputs=[text_input] | |
| ) | |
| # Toggle logic | |
| def toggle_comparison(is_comparing): | |
| return ( | |
| gr.Group(visible=not is_comparing), # single_model_group | |
| gr.Group(visible=is_comparing), # comparison_group | |
| gr.Group(visible=not is_comparing), # single_output_group | |
| gr.Group(visible=is_comparing), # comparison_output_group | |
| ) | |
| comparison_toggle.change( | |
| fn=toggle_comparison, | |
| inputs=[comparison_toggle], | |
| outputs=[single_model_group, comparison_group, single_output_group, comparison_output_group] | |
| ) | |
| # Single model run | |
| def run_single(model_name, input_text): | |
| scores, latency = run_model(model_name, input_text) | |
| return scores, latency | |
| # Comparison run | |
| def run_comparison_mode(model_1, model_2, input_text): | |
| scores_1, scores_2, latency_1, latency_2 = run_comparison(model_1, model_2, input_text) | |
| return scores_1, latency_1, scores_2, latency_2 | |
| def get_run_fn(is_comparing): | |
| if is_comparing: | |
| return run_comparison_mode, [model_dropdown_1, model_dropdown_2, text_input], [label_output_1, latency_output_1, label_output_2, latency_output_2] | |
| else: | |
| return run_single, [model_dropdown, text_input], [label_output, latency_output] | |
| run_button.click( | |
| fn=lambda comp, m1, m2, m_single, txt: ( | |
| ({}, "", *run_comparison_mode(m1, m2, txt)) if comp else (*run_single(m_single, txt), {}, "", {}, "") | |
| ), | |
| inputs=[comparison_toggle, model_dropdown_1, model_dropdown_2, model_dropdown, text_input], | |
| outputs=[label_output, latency_output, label_output_1, latency_output_1, label_output_2, latency_output_2], | |
| api_name="predict" | |
| ) | |
| def clear_all(): | |
| return "", {}, "", {}, "", {}, "" | |
| clear_btn.click( | |
| fn=clear_all, | |
| outputs=[text_input, label_output, latency_output, label_output_1, latency_output_1, label_output_2, latency_output_2] | |
| ) | |
| demo.launch(theme=gr.themes.Default(), css=custom_css) | |