waterfall 0.1.1__py3-none-any.whl → 0.1.4__py3-none-any.whl

waterfall/WatermarkerBase.py

@@ -141,6 +141,7 @@ class Watermarker:
                 do_sample=do_sample,
                 logits_processor=logits_processor,
                 pad_token_id=self.tokenizer.eos_token_id,
+                tokenizer=self.tokenizer,
                 **kwargs
                 )
         output = output[:,tokd_input["input_ids"].shape[-1]:].cpu()

waterfall/watermark.py

@@ -1,10 +1,12 @@
 import argparse
 import logging
 import os
+import gc
 import torch
 from typing import List, Literal, Optional, Tuple
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers.modeling_utils import PreTrainedModel
 from sentence_transformers import SentenceTransformer
 from tqdm.auto import tqdm
 
@@ -21,6 +23,8 @@ PROMPT = (
 )
 PRE_PARAPHRASED = "Here is a paraphrased version of the text while preserving the semantic similarity:\n\n"
 
+waterfall_cached_watermarking_model = None  # Global variable to cache the watermarking model
+
 def detect_gpu() -> str:
     """
     Use torch to detect if MPS, CUDA, or neither (default CPU)
@@ -42,9 +46,10 @@ def watermark(
     sts_model: SentenceTransformer,
     num_beam_groups: int = 4,
     beams_per_group: int = 2,
-    STS_scale:float = 2.0,
+    STS_scale: float = 2.0,
     diversity_penalty: float = 0.5,
     max_new_tokens: Optional[int] = None,
+    **kwargs
 ) -> str:
     paraphrasing_prompt = watermarker.tokenizer.apply_chat_template(
         [
@@ -61,6 +66,7 @@ def watermark(
         num_beam_groups = num_beam_groups,
         num_return_sequences = num_beam_groups * beams_per_group,
         diversity_penalty = diversity_penalty,
+        **kwargs,
         )
 
     # Select best paraphrasing based on q_score and semantic similarity
@@ -140,6 +146,7 @@ def watermark_texts(
     diversity_penalty: float = 0.5,
     STS_scale:float = 2.0,
     use_tqdm: bool = False,
+    stop_at_double_newline: bool = True,    # if True, will stop generation at the first double newline. Prevent repeated paraphrasing of the same text.
 ) -> List[str]:
     if watermark_fn == 'fourier':
         watermarkingFnClass = WatermarkingFnFourier
@@ -150,11 +157,25 @@ def watermark_texts(
 
     if watermarker is None:
         assert model_path is not None, "model_path must be provided if watermarker is not passed"
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            torch_dtype=torch_dtype,
-            device_map=device,
-            )
+        global waterfall_cached_watermarking_model
+
+        if isinstance(waterfall_cached_watermarking_model, PreTrainedModel) and waterfall_cached_watermarking_model.name_or_path != model_path:
+            device = waterfall_cached_watermarking_model.device.type
+            del waterfall_cached_watermarking_model
+            gc.collect()
+            if device == "cuda":
+                torch.cuda.empty_cache()
+            elif device == "mps":
+                torch.mps.empty_cache()
+            waterfall_cached_watermarking_model = None
+
+        if waterfall_cached_watermarking_model is None:
+            waterfall_cached_watermarking_model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=torch_dtype,
+                device_map=device,
+                )
+        model = waterfall_cached_watermarking_model
         tokenizer = AutoTokenizer.from_pretrained(model_path)
 
         watermarker = Watermarker(tokenizer=tokenizer, model=model, id=id, kappa=kappa, k_p=k_p, watermarkingFnClass=watermarkingFnClass)
@@ -173,6 +194,9 @@ def watermark_texts(
     T_ws = []
 
     for T_o in tqdm(T_os, desc="Watermarking texts",  disable=not use_tqdm):
+        if stop_at_double_newline and "\n\n" in T_o:
+            logging.warning("Text contains \\n\\n and stop_at_double_newline is set to True, replacing all \\n\\n in text.")
+            T_o = T_o.replace("\n\n", " ")  # replace double newlines with space
         T_w = watermark(
             T_o,
             watermarker = watermarker,
@@ -181,6 +205,7 @@ def watermark_texts(
             beams_per_group = beams_per_group,
             diversity_penalty = diversity_penalty,
             STS_scale = STS_scale,
+            stop_strings=["\n\n"] if stop_at_double_newline else None,
             )
         T_ws.append(T_w)
 

{waterfall-0.1.1.dist-info → waterfall-0.1.4.dist-info}/METADATA

@@ -1,23 +1,22 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.1.1
+Version: 0.1.4
 Summary: Scalable Framework for Robust Text Watermarking and Provenance for LLMs
-Author-email: Xinyuan Niu <aperture@outlook.sg>
-License-Expression: Apache-2.0
 Project-URL: Homepage, https://github.com/aoi3142/Waterfall
 Project-URL: Issues, https://github.com/aoi3142/Waterfall/issues
-Classifier: Programming Language :: Python :: 3
+Author-email: Xinyuan Niu <aperture@outlook.sg>
+License-Expression: Apache-2.0
+License-File: LICENSE
 Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
 Requires-Dist: accelerate>=0.29.0
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: scipy>=1.13.0
 Requires-Dist: sentence-transformers>=3.0.0
 Requires-Dist: torch>=2.3.0
 Requires-Dist: transformers>=4.43.1
-Dynamic: license-file
+Description-Content-Type: text/markdown
 
 # Waterfall: Scalable Framework for Robust Text Watermarking and Provenance for LLMs [EMNLP 2024 Main Long]
 Gregory Kang Ruey Lau*, Xinyuan Niu*, Hieu Dao, Jiangwei Chen, Chuan-Sheng Foo, Bryan Kian Hsiang Low
@@ -53,7 +52,7 @@ Protecting intellectual property (IP) of text such as articles and code is incre
 
 5. A token is sampled from the perturbed logits $\check{L}$ and is appended to the watermarked text.
 
-6. Append the generated token to the prompt and continue autoregressive generation (steps 1-5) until the eos token.
+6. Append the generated token to the prompt and continue autoregressive generation (steps 1-5) until the `eos` token.
 
 # Verification of un/watermarked text
 
@@ -77,6 +76,13 @@ Protecting intellectual property (IP) of text such as articles and code is incre
 
 # Using our code
 
+Install our package using `pip`
+```sh
+pip install waterfall
+```
+
+## Alternative installation from source
+
 [Optional]
 If using `conda` (or other pkg managers), it is highly advisable to create a new environment
 
@@ -98,6 +104,8 @@ Use the command `waterfall_demo` to watermark a piece of text, and then verify t
 waterfall_demo
 ```
 
+\* Ensure that your device (`cuda`/`cpu`/`mps`) has enough memory to load the model and perform inference (~18GB+ for default Llama 3.1 8B model)
+
 Additional arguments
 ```sh
 waterfall_demo \
@@ -107,9 +115,11 @@ waterfall_demo \
   --kappa        2                                `# Watermark strength`          \
   --model        meta-llama/Llama-3.1-8B-Instruct `# Paraphrasing LLM`            \
   --watermark_fn fourier                          `# fourier/square watermark`    \
-  --device       cuda                             `# Use cuda/cpu`
+  --device       cuda                             `# Use cuda/cpu/mps`
 ```
 
+\* By default, `--device` automatically selects among `cuda`/`cpu`/`mps` if not set
+
 ## Using our code to watermark and verify
 
 To watermark texts
@@ -134,7 +144,7 @@ test_texts = ["...", "..."]   # Suspected texts to verify
 watermark_strength = verify_texts(test_texts, id)[0]   # np array of floats
 ```
 
-# Code structure
+## Code structure
 
 - `watermark.py`              : Sample watermarking script used by with `watermark_demo` command, includes beam search and other optimizations
 - `WatermarkerBase.py`        : Underlying generation and verification code provided by `Watermarker` class

waterfall-0.1.4.dist-info/RECORD

@@ -0,0 +1,12 @@
+waterfall/WatermarkerBase.py,sha256=AyScrZz3hdjikvz5Fm4-B4acDz46i5wDFwCBg6Fp-vY,12947
+waterfall/WatermarkingFn.py,sha256=-b-kGRdL0a7eKRqJmcHPAR_rCjxQYnsg1Ne6bTwBc1I,1931
+waterfall/WatermarkingFnFourier.py,sha256=QYayAQYwi1dQkDIyqmvhU568VhrVYTVy47HkI8F8SZs,1358
+waterfall/WatermarkingFnSquare.py,sha256=2PAO05DdKT02npo7GDf_82D520nP7kGAWK6H4E4JMt4,1638
+waterfall/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+waterfall/permute.py,sha256=RwxOHFhx_VSOhhFwy5s79YgwTUBkfW2-LCCXYR3VT2o,2582
+waterfall/watermark.py,sha256=h7e1z8vWTUAKxCcQsJ2Jkx_1ZL-ug2dEDs5FzWcYfCs,13332
+waterfall-0.1.4.dist-info/METADATA,sha256=3hBQwb1JyrTWrayLCPFxXVlTpPjuE-ukPstW5F9F9rg,8715
+waterfall-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.1.4.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.1.4.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.1.4.dist-info/RECORD,,

{waterfall-0.1.1.dist-info → waterfall-0.1.4.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any
-

waterfall-0.1.1.dist-info/RECORD

@@ -1,13 +0,0 @@
-waterfall/WatermarkerBase.py,sha256=ou78I1XisalHbJLqyST6ryuLjtkFnY7Y60fUKdIwLy4,12905
-waterfall/WatermarkingFn.py,sha256=-b-kGRdL0a7eKRqJmcHPAR_rCjxQYnsg1Ne6bTwBc1I,1931
-waterfall/WatermarkingFnFourier.py,sha256=QYayAQYwi1dQkDIyqmvhU568VhrVYTVy47HkI8F8SZs,1358
-waterfall/WatermarkingFnSquare.py,sha256=2PAO05DdKT02npo7GDf_82D520nP7kGAWK6H4E4JMt4,1638
-waterfall/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-waterfall/permute.py,sha256=RwxOHFhx_VSOhhFwy5s79YgwTUBkfW2-LCCXYR3VT2o,2582
-waterfall/watermark.py,sha256=whiNhPwWNNIZwXMH6r7QzEE3A7Niq2Ro9elA1iSRoxI,11952
-waterfall-0.1.1.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.1.1.dist-info/METADATA,sha256=Ik8I-yLPuHSWdGsrSj7YgTCDJ0uTbfbV8FDvKYlPQ6M,8392
-waterfall-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-waterfall-0.1.1.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.1.1.dist-info/top_level.txt,sha256=5rTgijeT9V5GRCwIDZmhjeZ4khgH1lmfhS9ZmdUUCKQ,10
-waterfall-0.1.1.dist-info/RECORD,,

waterfall-0.1.1.dist-info/top_level.txt

@@ -1 +0,0 @@
-waterfall