now using hf's implementation for SynthID watermarking

detector_benchmark/create_dataset.py

@@ -6,6 +6,7 @@
     ElectraForSequenceClassification,
     ElectraTokenizer,
     AutoConfig,
+    SynthIDTextWatermarkingConfig,
 )
 import torch
 import argparse
@@ -62,12 +63,29 @@ def choose_watermarking_scheme(cfg: DictConfig, watermarking_scheme_name: str, g
 
     algorithm_config = cfg.watermark
 
-    watermarking_scheme = AutoWatermark.load(
-        watermarking_scheme_name,
-        algorithm_config=algorithm_config,
-        gen_model=gen,
-        model_config=model_config,
-    )
+    # temporary band-aid fix for SynthID while we try to find a better solution
+    if watermarking_scheme_name == "SynthID":
+
+        # we have to pass these parameters to the watermarking scheme
+        # but the true parameters are handled with the hydra config
+        watermarking_config = SynthIDTextWatermarkingConfig(
+            keys=[654, 400, 836, 123, 340, 443, 597, 160, 57, ...],
+            ngram_len=5,
+        )
+
+        for k, v in algorithm_config.items():
+            setattr(watermarking_config, k, v)
+
+        watermarking_scheme = watermarking_config
+
+    else:
+
+        watermarking_scheme = AutoWatermark.load(
+            watermarking_scheme_name,
+            algorithm_config=algorithm_config,
+            gen_model=gen,
+            model_config=model_config,
+        )
 
     return watermarking_scheme
 

detector_benchmark/generation/generator.py

@@ -9,6 +9,7 @@
     AutoModelForCausalLM,
     LogitsProcessor,
     LogitsProcessorList,
+    SynthIDTextWatermarkingConfig,
 )
 from ..utils.configs import ModelConfig
 from ..watermark.auto_watermark import AutoWatermark
@@ -68,12 +69,6 @@ def forward(
         outputs_list = []
         for i in tqdm(range(0, len(samples), batch_size), desc="Generating text"):
 
-            # specific for SynthID watermarking scheme, we need to reset the state
-            if watermarking_scheme is not None and hasattr(
-                watermarking_scheme.logits_processor, "state"
-            ):
-                watermarking_scheme.logits_processor.state = None
-
             batch_samples = samples[i : i + batch_size]
             encoding = self.tokenizer.batch_encode_plus(
                 batch_samples, return_tensors="pt", padding=True, truncation=True
@@ -83,8 +78,18 @@ def forward(
             with torch.no_grad():
                 if watermarking_scheme is not None:
 
+                    # special case for SynthID
+                    if isinstance(watermarking_scheme, SynthIDTextWatermarkingConfig):
+
+                        output_ids = self.generator.generate(
+                            input_ids,
+                            pad_token_id=self.tokenizer.pad_token_id,
+                            watermarking_config=watermarking_scheme,
+                            **self.gen_params
+                        )
+
                     # if the watermarking scheme has a logits processor, use it
-                    if hasattr(watermarking_scheme, "logits_processor"):
+                    elif hasattr(watermarking_scheme, "logits_processor"):
                         output_ids = self.generator.generate(
                             input_ids,
                             pad_token_id=self.tokenizer.pad_token_id,

detector_benchmark/watermark/synth_id/synth_id.py

@@ -58,10 +58,10 @@ def __init__(self, algorithm_config: dict, gen_model, model_config, *args, **kwa
         self.vocab_size = self.generation_tokenizer.vocab_size
         self.device = model_config.device
         self.gen_kwargs = model_config.gen_params
-        # self.top_k = model_config.gen_params["top_k"]
-        self.top_k = -1
-        # self.temperature = model_config.gen_params["temperature"]
-        self.temperature = 0.7
+        self.top_k = model_config.gen_params["top_k"]
+        # self.top_k = -1
+        self.temperature = model_config.gen_params["temperature"]
+        # self.temperature = 0.7
 
 
 class SynthIDUtils:

requirements.txt

@@ -137,13 +137,13 @@ tempdir==0.7.1
 termcolor==2.4.0
 threadpoolctl==3.3.0
 tinycss2==1.3.0
-tokenizers==0.19.1
+tokenizers==0.20.3
 torch>=2.2.1
 tornado==6.4
 tqdm==4.66.2
 tqdm-multiprocess==0.0.11
 traitlets==5.14.3
-transformers==4.44.2
+transformers==4.46.2
 typepy==1.3.2
 typing_extensions
 tyro==0.8.3