Update app.py

app.py

@@ -4,7 +4,7 @@ import json
 
 class SynthIDApp:
     def __init__(self):
-        self.api_url = "https://api-inference.huggingface.co/models/google/gemma-2b"
+        self.api_url = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
         self.headers = None
         self.WATERMARK_KEYS = [654, 400, 836, 123, 340, 443, 597, 160, 57, 789]
     
@@ -34,12 +34,13 @@ 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]"
             params = {
-                "inputs": text,
+                "inputs": prompt,
                 "parameters": {
-                    "max_new_tokens": 0,  # Don't generate new tokens
+                    "max_new_tokens": len(text.split()) * 2,  # Allow some flexibility
                     "do_sample": True,
-                    "temperature": 0.7,
+                    "temperature": 0.1,  # Lower temperature for more faithful reproduction
                     "top_p": 0.9,
                     "watermarking_config": {
                         "keys": self.WATERMARK_KEYS,
@@ -72,6 +73,11 @@ class SynthIDApp:
                     return text, f"API Error: {result[0]['error']}"
                 
                 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()
+                
                 if not watermarked_text:
                     return text, "Error: No watermarked text generated"
                 
@@ -112,7 +118,7 @@ app_instance = SynthIDApp()
 
 with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
     gr.Markdown("# SynthID Text Watermarking Tool")
-    gr.Markdown("Using Hugging Face Inference API for faster processing")
+    gr.Markdown("Using Mistral-7B-Instruct-v0.2 with Hugging Face Inference API")
     
     # Login section
     with gr.Row():