| import gradio as gr | |
| from models.watermark_faster import watermark_model | |
| import pdb | |
| from options import get_parser_main_model | |
| opts = get_parser_main_model().parse_args() | |
| model = watermark_model(language='English', mode=opts.mode, tau_word=0.8, lamda=0.83) | |
| def watermark_embed_demo(raw): | |
| watermarked_text = model.embed(raw) | |
| return watermarked_text | |
| def watermark_extract(raw): | |
| is_watermark, p_value, n, ones, z_value = model.watermark_detector_fast(raw) | |
| confidence = (1 - p_value) * 100 | |
| return f"{confidence:.2f}%" | |
| def precise_watermark_detect(raw): | |
| is_watermark, p_value, n, ones, z_value = model.watermark_detector_precise(raw) | |
| confidence = (1 - p_value) * 100 | |
| return f"{confidence:.2f}%" | |
| demo = gr.Blocks() | |
| with demo: | |
| with gr.Column(): | |
| gr.Markdown("# Watermarking Text Generated by Black-Box Language Models") | |
| inputs = gr.TextArea(label="Input text", placeholder="Copy your text here...") | |
| output = gr.Textbox(label="Watermarked Text") | |
| analysis_button = gr.Button("Inject Watermark") | |
| inputs_embed = [inputs] | |
| analysis_button.click(fn=watermark_embed_demo, inputs=inputs_embed, outputs=output) | |
| inputs_w = gr.TextArea(label="Text to Analyze", placeholder="Copy your watermarked text here...") | |
| mode = gr.Dropdown( | |
| label="Detection Mode", choices=["Fast", "Precise"], default="Fast" | |
| ) | |
| output_detect = gr.Textbox(label="Confidence (the likelihood of the text containing a watermark)") | |
| detect_button = gr.Button("Detect") | |
| def detect_watermark(inputs_w, mode): | |
| if mode == "Fast": | |
| return watermark_extract(inputs_w) | |
| else: | |
| return precise_watermark_detect(inputs_w) | |
| detect_button.click(fn=detect_watermark, inputs=[inputs_w, mode], outputs=output_detect) | |
| if __name__ == "__main__": | |
| gr.close_all() | |
| demo.title = "Watermarking Text Generated by Black-Box Language Models" | |
| demo.launch() |