Update app.py

app.py

@@ -34,21 +34,14 @@ class SynthIDApp:
         try:
             # Prepare the API request parameters
             # Prepare the API request parameters for watermarking
-            prompt = f"<s>[INST] Apply watermark to the following text WITHOUT adding any extra words or changing the meaning: {text} [/INST]"
-            
-            # Calculate exact token limits
-            input_length = len(text.split())
+            prompt = f"<s>[INST] Return the exact same text, with watermark applied: {text} [/INST]"
             
             params = {
                 "inputs": prompt,
                 "parameters": {
-                    "max_new_tokens": input_length + 2,    # Very strict limit
-                    "min_new_tokens": input_length - 2,    # Allow tiny variations
-                    "do_sample": True,
-                    "temperature": 0.01,                   # Almost deterministic
-                    "top_p": 0.99,
-                    "repetition_penalty": 1.5,            # Strongly prevent repetitions
-                    "stop": ["[/INST]", "\n", "."],      # Stop at natural boundaries
+                    "return_full_text": True,
+                    "do_sample": False,     # Deterministic generation
+                    "temperature": 0.01,    # Almost deterministic
                     "watermarking_config": {
                         "keys": self.WATERMARK_KEYS,
                         "ngram_len": int(ngram_len)
@@ -79,28 +72,31 @@ class SynthIDApp:
                 if 'error' in result[0]:
                     return text, f"API Error: {result[0]['error']}"
                 
-                watermarked_text = result[0].get('generated_text', '').strip()
+                generated_text = result[0].get('generated_text', '').strip()
                 
-                # Clean up Mistral's response format and extract only the relevant part
-                parts = watermarked_text.split("[/INST]")
-                if len(parts) > 1:
-                    watermarked_text = parts[-1].strip()
+                # Extract only the response part after the instruction
+                try:
+                    # First try splitting on [/INST]
+                    parts = generated_text.split("[/INST]")
+                    if len(parts) > 1:
+                        watermarked_text = parts[-1].strip()
+                    else:
+                        # If no [/INST], try finding the original text and take what follows
+                        idx = generated_text.find(text)
+                        if idx != -1:
+                            watermarked_text = generated_text[idx + len(text):].strip()
+                        else:
+                            # If all else fails, take the whole text
+                            watermarked_text = generated_text
+                except Exception as e:
+                    return text, f"Error processing response: {str(e)}"
                 
-                # Remove any trailing periods or extra spaces
-                watermarked_text = watermarked_text.rstrip('.')
-                while '  ' in watermarked_text:
-                    watermarked_text = watermarked_text.replace('  ', ' ')
+                # Clean up the text
+                watermarked_text = watermarked_text.strip(' .')
                 
                 if not watermarked_text:
                     return text, "Error: No watermarked text generated"
                 
-                # Very strict length validation
-                input_words = len(text.split())
-                output_words = len(watermarked_text.split())
-                
-                if output_words != input_words:
-                    return text, f"Error: Length mismatch ({output_words} vs {input_words} words)"
-                
                 # Add back the period if the original had one
                 if text.strip().endswith('.'):
                     watermarked_text += '.'