waterfall 0.2.8__py3-none-any.whl → 0.2.9__py3-none-any.whl

waterfall/WatermarkerBase.py

@@ -15,7 +15,7 @@ from tqdm import tqdm
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers.modeling_utils import PreTrainedModel
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase, BatchEncoding
-from transformers.generation.logits_process import LogitsProcessor, TopKLogitsWarper, TopPLogitsWarper
+from transformers.generation.logits_process import LogitsProcessor, TopKLogitsWarper, TopPLogitsWarper, TemperatureLogitsWarper
 from transformers.generation.configuration_utils import GenerationConfig
 
 from waterfall.permute import Permute
@@ -265,11 +265,23 @@ class Watermarker:
             generation_config: GenerationConfig = kwargs["generation_config"]
             top_k = generation_config.top_k
             top_p = generation_config.top_p
-            generation_config.update(top_p=1.0)
+            temperature = generation_config.temperature
+            num_beams = generation_config.num_beams
+            diversity_penalty = generation_config.diversity_penalty
+            if num_beams <= 1:
+                diversity_penalty = None
+            generation_config.update(top_p=1.0, temperature=None, diversity_penalty=diversity_penalty)
         else:
             top_k = kwargs.pop("top_k", None)
             top_p = kwargs.pop("top_p", None)
-
+            temperature = kwargs.pop("temperature", 1.0)
+            num_beams = kwargs.pop("num_beams", 1)
+            diversity_penalty = kwargs.pop("diversity_penalty", None)
+            if num_beams <= 1:
+                kwargs["diversity_penalty"] = None
+
+        if num_beams > 1 and temperature != 1.0:
+            logits_processor.append(TemperatureLogitsWarper(float(temperature)))
         if top_k is not None and top_k != 0:
             logits_processor.append(TopKLogitsWarper(top_k))
         if top_p is not None and top_p < 1.0:

waterfall/watermark.py

@@ -39,6 +39,17 @@ def detect_gpu() -> str:
     else:
         return 'cpu'
 
+def del_cached_model():
+    global waterfall_cached_watermarking_model
+    if isinstance(waterfall_cached_watermarking_model, PreTrainedModel):
+        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()
+
 def watermark_texts(
     T_os: List[str],
     id: Optional[int] = None,
@@ -60,6 +71,7 @@ def watermark_texts(
     beams_per_group: int = 2,
     diversity_penalty: float = 0.5,
     stop_at_double_newline: bool = True,    # if True, will stop generation at the first double newline. Prevent repeated paraphrasing of the same text.
+    **kwargs,
 ) -> List[str]:
     if watermark_fn == 'fourier':
         watermarkingFnClass = WatermarkingFnFourier
@@ -72,16 +84,8 @@ def watermark_texts(
     if watermarker is None:
         assert model_path is not None, "model_path must be provided if watermarker is not passed"
         assert id is not None, "id must be provided if watermarker is not passed"
-        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()
+            del_cached_model()
 
         if waterfall_cached_watermarking_model is None:
             model = model_path
@@ -91,7 +95,10 @@ def watermark_texts(
         watermarker = Watermarker(model=model, id=id, kappa=kappa, k_p=k_p, watermarkingFnClass=watermarkingFnClass)
     else:
         device = watermarker.model.device.type
-        id = watermarker.id
+        if id is not None:
+            watermarker.set_id(id)
+        else:
+            id = watermarker.id
     waterfall_cached_watermarking_model = watermarker.model
 
     # Check if sts model is loaded
@@ -122,8 +129,8 @@ def watermark_texts(
         assert (do_sample and temperature is not None and top_p is not None and num_beam_groups == 1 and beams_per_group == 1), \
            "do_sample=True requires temperature, top_p, num_beam_groups=1 and beams_per_group=1"
     else:   # Using beam search
-        assert (not do_sample and temperature is None and top_p is None and num_beam_groups >= 1 and beams_per_group >= 1), \
-           "do_sample=False requires temperature=None, top_p=None, num_beam_groups>=1 and beams_per_group>=1"
+        assert (not do_sample and num_beam_groups >= 1 and beams_per_group >= 1), \
+           "do_sample=False requires num_beam_groups>=1 and beams_per_group>=1"
 
     eos_token_id = watermarker.tokenizer.eos_token_id
     # add "\n\n" tokens to eos_token_id list
@@ -144,6 +151,7 @@ def watermark_texts(
         diversity_penalty=diversity_penalty,
         eos_token_id=eos_token_id,
         num_return_sequences=num_beam_groups * beams_per_group,
+        **kwargs
     )
 
     watermarked = watermarker.generate(
@@ -152,6 +160,7 @@ def watermark_texts(
         return_scores=True,
         use_tqdm=use_tqdm,
         generation_config=generation_config,
+        use_model_defaults=False,
     )
     T_ws = watermarked["text"]
     # Reshape T_ws to Queries X Beams
@@ -179,10 +188,7 @@ def verify_texts(texts: List[str], id: int,
         assert model_path is not None, "model_path must be provided if watermarker is not passed"
         watermarker = Watermarker(tokenizer=model_path)
 
-    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=return_extracted_k_p)  # 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]

{waterfall-0.2.8.dist-info → waterfall-0.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.2.8
+Version: 0.2.9
 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
@@ -76,7 +76,7 @@ Protecting intellectual property (IP) of text such as articles and code is incre
 
 # Using our code
 
-Install our package using `pip`
+Install our [PyPI package](https://pypi.org/project/waterfall/) using `pip`
 ```sh
 pip install waterfall
 ```

{waterfall-0.2.8.dist-info → waterfall-0.2.9.dist-info}/RECORD

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=y7rJtP4Qf8GdJNdlZriX-JLClKxJRGiCtIkBtDoIqZw,21300
+waterfall/WatermarkerBase.py,sha256=ZyH1AWlXl1ak9lCnhtBdsd4IXsHKFGNnx_yYJFnf_zw,22018
 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=uYKdmn4pGvjB6hInInLGxFIF6vt507lqJ_qe-ST1PFE,2783
-waterfall/watermark.py,sha256=IbH5r3oqjtKztDVryfDTr_NDn-CLZHow0S8nAEtZmdc,14420
-waterfall-0.2.8.dist-info/METADATA,sha256=QO_d3epwMllDm1vwvcOwwKL0ho5T6X4NEXXd7YHfiso,8723
-waterfall-0.2.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-waterfall-0.2.8.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.2.8.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.2.8.dist-info/RECORD,,
+waterfall/watermark.py,sha256=avyQIFJBhqu_q_ZBp0-RWvAOIJmzJvVisbiIca2GPyA,14536
+waterfall-0.2.9.dist-info/METADATA,sha256=9sFs2kA_fj5CUD2UBt9uEuB1UTgn-ITJWUFWkksEqzE,8767
+waterfall-0.2.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.2.9.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.2.9.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.2.9.dist-info/RECORD,,