Update app.py

app.py

@@ -1,28 +1,31 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import requests
 from transformers import SynthIDTextWatermarkingConfig
 
 class SynthIDApp:
     def __init__(self):
-        self.client = None
+        self.api_url = "https://api-inference.huggingface.co/models/google/gemma-2b"
+        self.headers = 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."""
+        """Initialize the API headers with authentication."""
         try:
-            # Initialize the inference client
-            self.client = InferenceClient(
-                model="google/gemma-2b",
-                token=hf_token
+            self.headers = {"Authorization": f"Bearer {hf_token}"}
+            
+            # Test the connection with a simple query
+            response = requests.post(
+                self.api_url,
+                headers=self.headers,
+                json={"inputs": "Test", "parameters": {"max_new_tokens": 1}}
             )
+            response.raise_for_status()
             
-            # Test the connection with a simple generation
-            _ = self.client.text_generation("Test", max_new_tokens=1)
-            return "Inference client initialized successfully!"
+            return "API connection initialized successfully!"
         except Exception as e:
-            self.client = None
-            return f"Error initializing client: {str(e)}"
+            self.headers = None
+            return f"Error initializing API: {str(e)}"
 
     def update_watermark_config(self, ngram_len):
         """Update the watermarking configuration with new ngram_len."""
@@ -36,32 +39,44 @@ class SynthIDApp:
             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."
+        """Apply SynthID watermark to input text using the inference API."""
+        if not self.headers:
+            return text, "Error: API 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,
-                "ngram_len": self.watermarking_config.ngram_len
+            # Prepare the API request parameters
+            params = {
+                "inputs": text,
+                "parameters": {
+                    "max_new_tokens": 100,
+                    "do_sample": True,
+                    "temperature": 0.7,
+                    "top_p": 0.9,
+                    "watermarking_config": {
+                        "keys": self.watermarking_config.keys,
+                        "ngram_len": self.watermarking_config.ngram_len
+                    }
+                }
             }
             
-            # Make the API call with watermarking config
-            response = self.client.text_generation(
-                text,
-                max_new_tokens=100,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                watermarking_config=watermark_dict,
-                return_full_text=False
+            # Make the API call
+            response = requests.post(
+                self.api_url,
+                headers=self.headers,
+                json=params
             )
+            response.raise_for_status()
             
-            watermarked_text = response
+            # Extract the generated text
+            result = response.json()
+            if isinstance(result, list) and len(result) > 0:
+                watermarked_text = result[0].get('generated_text', text)
+            else:
+                watermarked_text = text
+                
             return watermarked_text, f"Watermark applied successfully! (ngram_len: {ngram_len})"
         except Exception as e:
             return text, f"Error applying watermark: {str(e)}"
@@ -89,7 +104,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 Endpoints for faster processing")
+    gr.Markdown("Using Hugging Face Inference API for faster processing")
     
     # Login section
     with gr.Row():
@@ -153,7 +168,7 @@ with gr.Blocks(title="SynthID Text Watermarking Tool") as app:
     3. Adjust the N-gram Length slider to control watermark characteristics
     
     ### Notes:
-    - This version uses Hugging Face's Inference Endpoints for faster processing
+    - This version uses Hugging Face's Inference API for faster processing
     - No model download required - everything runs in the cloud
     - The watermark is designed to be imperceptible to humans
     - This demo only implements watermark application