waterfall 0.1.3__py3-none-any.whl → 0.1.5__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
@@ -75,7 +81,8 @@ def watermark(
 def verify_texts(texts: List[str], id: int, 
                      watermarker: Optional[Watermarker] = None, 
                      k_p: Optional[int] = None, 
-                     model_path: Optional[str] = "meta-llama/Llama-3.1-8B-Instruct"
+                     model_path: Optional[str] = "meta-llama/Llama-3.1-8B-Instruct",
+                     return_extracted_k_p: bool = False
                      ) -> Tuple[float,float]:
     """Returns the q_score and extracted k_p"""
 
@@ -87,7 +94,11 @@ def verify_texts(texts: List[str], id: int,
     if k_p is None:
         k_p = watermarker.k_p
 
-    verify_results = watermarker.verify(texts, id=[id], k_p=[k_p], return_extracted_k_p=True)  # results are [text x id x k_p]
+    verify_results = watermarker.verify(texts, id=[id], k_p=[k_p], return_extracted_k_p=return_extracted_k_p)  # results are [text x id x k_p]
+
+    if not return_extracted_k_p:
+        return verify_results[:,0,0]
+
     q_score = verify_results["q_score"]
     k_p_extracted = verify_results["k_p_extracted"]
 
@@ -140,6 +151,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 +162,24 @@ 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
+            waterfall_cached_watermarking_model = None
+            gc.collect()
+            if device == "cuda":
+                torch.cuda.empty_cache()
+            elif device == "mps":
+                torch.mps.empty_cache()
+
+        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 +198,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 +209,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)
 
@@ -283,7 +312,7 @@ def main():
         )
 
     # watermarker = Watermarker(tokenizer=tokenizer, model=None, id=id, k_p=k_p, watermarkingFnClass=watermarkingFnClass)   # If only verifying the watermark, do not need to instantiate the model
-    q_scores, extracted_k_ps = verify_texts(T_os + T_ws, id, watermarker, k_p=k_p)
+    q_scores, extracted_k_ps = verify_texts(T_os + T_ws, id, watermarker, k_p=k_p, return_extracted_k_p=True)
 
     for i in range(len(T_os)):
         # Handle the case where this is being run

{waterfall-0.1.3.dist-info → waterfall-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.1.3
+Version: 0.1.5
 Summary: Scalable Framework for Robust Text Watermarking and Provenance for LLMs
 Project-URL: Homepage, https://github.com/aoi3142/Waterfall
 Project-URL: Issues, https://github.com/aoi3142/Waterfall/issues
@@ -141,7 +141,7 @@ from waterfall.watermark import verify_texts
 id = 1                        # specify your watermarking ID
 test_texts = ["...", "..."]   # Suspected texts to verify
 
-watermark_strength = verify_texts(test_texts, id)[0]   # np array of floats
+watermark_strength = verify_texts(test_texts, id)     # np array of floats
 ```
 
 ## Code structure

{waterfall-0.1.3.dist-info → waterfall-0.1.5.dist-info}/RECORD

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=ou78I1XisalHbJLqyST6ryuLjtkFnY7Y60fUKdIwLy4,12905
+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=whiNhPwWNNIZwXMH6r7QzEE3A7Niq2Ro9elA1iSRoxI,11952
-waterfall-0.1.3.dist-info/METADATA,sha256=hnX9Vq2zjuq_u9VxJcP9OQ5U2RHUtTGnTQu6iwUJhXM,8715
-waterfall-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-waterfall-0.1.3.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.1.3.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.1.3.dist-info/RECORD,,
+waterfall/watermark.py,sha256=d93C_DjnO2flLHZkRAC_L749CcUjIIGT9yPR91KkjOk,13452
+waterfall-0.1.5.dist-info/METADATA,sha256=aCbjPA6TZwVNcgrX9m0YDsaLcamx7-dayxKkjJmGUcA,8714
+waterfall-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.1.5.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.1.5.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.1.5.dist-info/RECORD,,