waterfall 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

waterfall/WatermarkerBase.py

@@ -27,13 +27,20 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # Check transformers version
 import transformers
 from packaging import version
-# Group beam search is shifted to transformers-community package in 4.57.0
-use_custom_group_beam_search = version.parse(transformers.__version__) >= version.parse("4.57.0")
+
+transformers_version = version.parse(transformers.__version__)
 # Set model loading kwargs based on transformers version
-if version.parse(transformers.__version__) >= version.parse("4.56.0"):
+if transformers_version >= version.parse("4.56.0"):
     model_from_pretrained_kwargs = {"dtype": "auto"}
 else:
     model_from_pretrained_kwargs = {"torch_dtype": torch.bfloat16}
+# Group beam search is shifted to transformers-community package in 4.57.0
+use_custom_group_beam_search = transformers_version >= version.parse("4.57.0")
+# BatchEncoding to() non_blocking added in 4.48.0
+if transformers_version >= version.parse("4.48.0"):
+    batch_encoding_to_kwargs = {"non_blocking": True}
+else:
+    batch_encoding_to_kwargs = {}
 
 class PerturbationProcessor(LogitsProcessor):
     def __init__(self,
@@ -316,7 +323,7 @@ class Watermarker:
             tokd_inputs_order = range(len(tokd_inputs))
         tokd_input_batches = []
         for i in range(0, len(tokd_inputs), max_batch_size):
-            batch = self.tokenizer.pad(tokd_inputs[i:i+max_batch_size], padding=True, padding_side="left").to(self.model.device, non_blocking=True)
+            batch = self.tokenizer.pad(tokd_inputs[i:i+max_batch_size], padding=True, padding_side="left").to(self.model.device, **batch_encoding_to_kwargs)
             tokd_input_batches.append(batch)
         torch.cuda.synchronize()
 

waterfall/watermark.py

@@ -18,10 +18,17 @@ from waterfall.WatermarkerBase import Watermarker
 # Check transformers version
 import transformers
 from packaging import version
-if version.parse(transformers.__version__) >= version.parse("4.56.0"):
+transformers_version = version.parse(transformers.__version__)
+if transformers_version >= version.parse("4.56.0"):
     model_from_pretrained_kwargs = {"dtype": "auto"}
 else:
     model_from_pretrained_kwargs = {"torch_dtype": torch.bfloat16}
+if transformers_version < version.parse("5.0.0") and transformers_version >= version.parse("4.50.0"):
+    additional_generation_config = {
+        "use_model_defaults": False,
+    }
+else:
+    additional_generation_config = {}
 
 PROMPT = (
     "Paraphrase the user provided text while preserving semantic similarity. "
@@ -168,7 +175,7 @@ def watermark_texts(
         return_scores=True,
         use_tqdm=use_tqdm,
         generation_config=generation_config,
-        use_model_defaults=False,
+        **additional_generation_config,
     )
     T_ws = watermarked["text"]
     # Reshape T_ws to Queries X Beams

{waterfall-0.2.12.dist-info → waterfall-0.2.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.2.12
+Version: 0.2.13
 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.2.12.dist-info → waterfall-0.2.13.dist-info}/RECORD

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=6O_S78dD3Jha2OkJK2u3euwCH93i-mTiYYGXosPDMig,22632
+waterfall/WatermarkerBase.py,sha256=H-tJ96WUihW30EFFnPn92pna4qQtyYjcWBlVVtY3oMM,22863
 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=fvscFoSbM51YUuDaOmrOKGvwXO25VMgGJKTfAeeKCaA,14817
-waterfall-0.2.12.dist-info/METADATA,sha256=TBoeAFK8qkG-jIRi-OeKq4GbtFTDQkaUKyKuFNeDQHo,8760
-waterfall-0.2.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-waterfall-0.2.12.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.2.12.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.2.12.dist-info/RECORD,,
+waterfall/watermark.py,sha256=Qe_NSNH2XL5ZCf069fa438NOpNsju3l4kr2GDoKbuVU,15093
+waterfall-0.2.13.dist-info/METADATA,sha256=VwO9mXTFEOoFxRASPt7qeZVCkMIbhH3_LkJ02yccOFM,8760
+waterfall-0.2.13.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+waterfall-0.2.13.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.2.13.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.2.13.dist-info/RECORD,,