File size: 3,434 Bytes
5ea9e86 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import gradio as gr
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
SynthIDTextWatermarkingConfig,
SynthIDTextBayesianDetector
)
# Initialize model and tokenizer
MODEL_NAME = "google/gemma-2b" # You can change this to your preferred model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
# Configure watermarking
WATERMARK_KEYS = [654, 400, 836, 123, 340, 443, 597, 160, 57, 789] # Example keys
watermarking_config = SynthIDTextWatermarkingConfig(
keys=WATERMARK_KEYS,
ngram_len=5
)
# Initialize detector
detector = SynthIDTextBayesianDetector(watermarking_config)
def apply_watermark(text):
"""Apply SynthID watermark to input text."""
try:
# Tokenize input
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
# Generate with watermark
with torch.no_grad():
outputs = model.generate(
**inputs,
watermarking_config=watermarking_config,
do_sample=True,
max_length=len(inputs["input_ids"][0]) + 100, # Add some extra tokens
pad_token_id=tokenizer.eos_token_id
)
# Decode output
watermarked_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return watermarked_text, "Watermark applied successfully!"
except Exception as e:
return text, f"Error applying watermark: {str(e)}"
def detect_watermark(text):
"""Detect if text contains SynthID watermark."""
try:
# Get detection score
score = detector.detect(text)
# Interpret results
threshold = 0.5 # You can adjust this threshold
is_watermarked = score > threshold
result = f"Watermark Detection Score: {score:.3f}\n"
result += f"Verdict: {'WATERMARK DETECTED' if is_watermarked else 'NO WATERMARK DETECTED'}"
return result
except Exception as e:
return f"Error detecting watermark: {str(e)}"
# Create Gradio interface
with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
gr.Markdown("# SynthID Text Watermarking Tool")
gr.Markdown("Apply and detect SynthID watermarks in text")
with gr.Tab("Apply Watermark"):
with gr.Row():
input_text = gr.Textbox(label="Input Text", lines=5)
output_text = gr.Textbox(label="Watermarked Text", lines=5)
status = gr.Textbox(label="Status")
apply_btn = gr.Button("Apply Watermark")
apply_btn.click(apply_watermark, inputs=[input_text], outputs=[output_text, status])
with gr.Tab("Detect Watermark"):
with gr.Row():
detect_input = gr.Textbox(label="Text to Check", lines=5)
detect_result = gr.Textbox(label="Detection Result", lines=3)
detect_btn = gr.Button("Detect Watermark")
detect_btn.click(detect_watermark, inputs=[detect_input], outputs=[detect_result])
gr.Markdown("""
### Notes:
- The watermark is designed to be imperceptible to humans but detectable by the classifier
- Detection scores above 0.5 indicate likely presence of a watermark
- The watermark is somewhat robust to minor text modifications but may not survive major changes
""")
# Launch the app
if __name__ == "__main__":
app.launch() |