waterfall 0.1.3__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.3.dist-info → waterfall-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.1.3
+Version: 0.1.4
 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

{waterfall-0.1.3.dist-info → waterfall-0.1.4.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=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,,