Update app.py

app.py

@@ -34,14 +34,16 @@ class SynthIDApp:
         try:
             # Prepare the API request parameters
             # Prepare the API request parameters for watermarking
-            prompt = f"<s>[INST] Rewrite this text exactly as is: {text} [/INST]"
+            prompt = f"<s>[INST] Apply watermark to this text: {text} [/INST]"
             params = {
                 "inputs": prompt,
                 "parameters": {
-                    "max_new_tokens": len(text.split()) * 2,  # Allow some flexibility
+                    "max_new_tokens": len(text.split()) + 10,  # Limit to input length plus some buffer
+                    "min_new_tokens": len(text.split()),       # At least return same length
                     "do_sample": True,
-                    "temperature": 0.1,  # Lower temperature for more faithful reproduction
-                    "top_p": 0.9,
+                    "temperature": 0.1,     # Low temperature for more faithful reproduction
+                    "top_p": 0.95,
+                    "repetition_penalty": 1.2,  # Prevent repetitions
                     "watermarking_config": {
                         "keys": self.WATERMARK_KEYS,
                         "ngram_len": int(ngram_len)
@@ -75,17 +77,19 @@ class SynthIDApp:
                 watermarked_text = result[0].get('generated_text', '').strip()
                 
                 # Clean up Mistral's response format
-                watermarked_text = watermarked_text.replace("<s>[INST]", "").replace("[/INST]", "")
-                watermarked_text = watermarked_text.replace("Rewrite this text exactly as is:", "").strip()
+                watermarked_text = watermarked_text.split("[/INST]")[-1].strip()
                 
                 if not watermarked_text:
                     return text, "Error: No watermarked text generated"
                 
-                # Add basic length validation
-                if len(watermarked_text) < len(text) * 0.5:
-                    return text, "Error: Generated text too short"
-                if len(watermarked_text) > len(text) * 1.5:
-                    return text, "Error: Generated text too long"
+                # More precise length validation
+                input_words = len(text.split())
+                output_words = len(watermarked_text.split())
+                
+                if output_words < input_words * 0.8:
+                    return text, f"Error: Generated text too short ({output_words} vs {input_words} words)"
+                if output_words > input_words * 1.2:
+                    return text, f"Error: Generated text too long ({output_words} vs {input_words} words)"
                 
                 return watermarked_text, f"Watermark applied successfully! (ngram_len: {ngram_len})"
             else:
@@ -138,7 +142,7 @@ with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
                     label="Input Text", 
                     lines=5,
                     placeholder="Enter text to watermark...",
-                    value="The quick brown fox jumps over the lazy dog."
+                    value="This is a test sentence for the watermarking system."
                 )
                 output_text = gr.Textbox(label="Watermarked Text", lines=5)
             with gr.Column(scale=1):