Update app.py

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 from transformers import SynthIDTextWatermarkingConfig
-import json
 
 class SynthIDApp:
     def __init__(self):
         self.client = None
         self.watermarking_config = None
+        self.WATERMARK_KEYS = [654, 400, 836, 123, 340, 443, 597, 160, 57, 789]
     
     def login(self, hf_token):
         """Initialize the inference client with authentication."""
@@ -17,27 +17,33 @@ class SynthIDApp:
                 token=hf_token
             )
             
-            # Configure watermarking
-            WATERMARK_KEYS = [654, 400, 836, 123, 340, 443, 597, 160, 57, 789]
-            self.watermarking_config = SynthIDTextWatermarkingConfig(
-                keys=WATERMARK_KEYS,
-                ngram_len=5
-            )
-            
-            # Test the connection
-            _ = self.client.token_count("Test")
+            # Test the connection with a simple generation
+            _ = self.client.text_generation("Test", max_new_tokens=1)
             return "Inference client initialized successfully!"
         except Exception as e:
             self.client = None
-            self.watermarking_config = None
             return f"Error initializing client: {str(e)}"
 
-    def apply_watermark(self, text):
+    def update_watermark_config(self, ngram_len):
+        """Update the watermarking configuration with new ngram_len."""
+        try:
+            self.watermarking_config = SynthIDTextWatermarkingConfig(
+                keys=self.WATERMARK_KEYS,
+                ngram_len=ngram_len
+            )
+            return f"Watermark config updated: ngram_len = {ngram_len}"
+        except Exception as e:
+            return f"Error updating config: {str(e)}"
+
+    def apply_watermark(self, text, ngram_len):
         """Apply SynthID watermark to input text using the inference endpoint."""
         if not self.client:
             return text, "Error: Client not initialized. Please login first."
             
         try:
+            # Update watermark config with current ngram_len
+            self.update_watermark_config(ngram_len)
+            
             # Convert watermarking config to dict for the API call
             watermark_dict = {
                 "keys": self.watermarking_config.keys,
@@ -56,7 +62,7 @@ class SynthIDApp:
             )
             
             watermarked_text = response
-            return watermarked_text, "Watermark applied successfully!"
+            return watermarked_text, f"Watermark applied successfully! (ngram_len: {ngram_len})"
         except Exception as e:
             return text, f"Error applying watermark: {str(e)}"
 
@@ -65,19 +71,12 @@ class SynthIDApp:
         try:
             total_words = len(text.split())
             avg_word_length = sum(len(word) for word in text.split()) / total_words if total_words > 0 else 0
-            
-            # Get token count if client is available
-            token_info = ""
-            if self.client:
-                try:
-                    token_count = self.client.token_count(text)
-                    token_info = f"\n- Token count: {token_count}"
-                except:
-                    pass
+            char_count = len(text)
             
             analysis = f"""Text Analysis:
+- Total characters: {char_count}
 - Total words: {total_words}
-- Average word length: {avg_word_length:.2f}{token_info}
+- Average word length: {avg_word_length:.2f}
 
 Note: This is a basic analysis. The official SynthID detector is not yet available in the public transformers package."""
             
@@ -105,15 +104,36 @@ with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
     
     with gr.Tab("Apply Watermark"):
         with gr.Row():
-            input_text = gr.Textbox(
-                label="Input Text", 
-                lines=5,
-                placeholder="Enter text to watermark..."
-            )
-            output_text = gr.Textbox(label="Watermarked Text", lines=5)
-            status = gr.Textbox(label="Status")
+            with gr.Column(scale=3):
+                input_text = gr.Textbox(
+                    label="Input Text", 
+                    lines=5,
+                    placeholder="Enter text to watermark..."
+                )
+                output_text = gr.Textbox(label="Watermarked Text", lines=5)
+            with gr.Column(scale=1):
+                ngram_len = gr.Slider(
+                    label="N-gram Length",
+                    minimum=2,
+                    maximum=5,
+                    step=1,
+                    value=5,
+                    info="Controls watermark detectability (2-5)"
+                )
+                status = gr.Textbox(label="Status")
+        
+        gr.Markdown("""
+        ### N-gram Length Parameter:
+        - Higher values (4-5): More detectable watermark, but more brittle to changes
+        - Lower values (2-3): More robust to changes, but harder to detect
+        - Default (5): Maximum detectability""")
+        
         apply_btn = gr.Button("Apply Watermark")
-        apply_btn.click(app_instance.apply_watermark, inputs=[input_text], outputs=[output_text, status])
+        apply_btn.click(
+            app_instance.apply_watermark, 
+            inputs=[input_text, ngram_len], 
+            outputs=[output_text, status]
+        )
     
     with gr.Tab("Analyze Text"):
         with gr.Row():
@@ -130,6 +150,7 @@ with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
     ### Instructions:
     1. Enter your Hugging Face token and click Login
     2. Once connected, you can use the tabs to apply watermarks or analyze text
+    3. Adjust the N-gram Length slider to control watermark characteristics
     
     ### Notes:
     - This version uses Hugging Face's Inference Endpoints for faster processing