Create app.py

app.py

@@ -0,0 +1,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()