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

waterfall/WatermarkerBase.py

@@ -228,6 +228,7 @@ class Watermarker:
             use_tqdm : bool = False,
             batched_generate : bool = True,
             discard_incomplete : bool = True,
+            logits_processor = [],
             **kwargs    # Other generate parameters
             ) -> List[str] | dict:  # Returns flattened list of query x beam
 
@@ -259,7 +260,6 @@ class Watermarker:
             squeezed_tokd_inputs.append(BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask}))
         tokd_inputs = squeezed_tokd_inputs
 
-        logits_processor = []
         # Ensure top_k and top_p happens before watermarking
         if "generation_config" in kwargs:
             generation_config: GenerationConfig = kwargs["generation_config"]
@@ -275,12 +275,12 @@ class Watermarker:
             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)
+            num_beams = kwargs.get("num_beams", 1)
+            diversity_penalty = kwargs.get("diversity_penalty", None)
             if num_beams <= 1:
                 kwargs["diversity_penalty"] = None
 
-        if num_beams > 1 and temperature != 1.0:
+        if num_beams > 1 and temperature is not None 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))
@@ -351,10 +351,6 @@ class Watermarker:
             decoded_output = [i.strip() for i in decoded_output]
             return_dict["text"] = decoded_output
 
-        if is_single:
-            for k, v in return_dict.items():
-                return_dict[k] = v[0]
-
         if return_text and len(return_dict) == 1:
             return decoded_output
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.2.9
+Version: 0.2.11
 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
@@ -15,7 +15,7 @@ Requires-Dist: numpy>=1.25.0
 Requires-Dist: scipy>=1.13.0
 Requires-Dist: sentence-transformers>=3.0.0
 Requires-Dist: torch>=2.3.0
-Requires-Dist: transformers<4.55.0,>=4.43.1
+Requires-Dist: transformers<4.57.0,>=4.43.1
 Description-Content-Type: text/markdown
 
 # Waterfall: Scalable Framework for Robust Text Watermarking and Provenance for LLMs [EMNLP 2024 Main Long]

waterfall-0.2.11.dist-info/RECORD

@@ -0,0 +1,12 @@
+waterfall/WatermarkerBase.py,sha256=pk2NB7J0oBLXcO0FIRBHllnSowbpeRNd9ZjvPuUOeeM,21945
+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=avyQIFJBhqu_q_ZBp0-RWvAOIJmzJvVisbiIca2GPyA,14536
+waterfall-0.2.11.dist-info/METADATA,sha256=Ttp-F0sjA31gppuF6dfo5ze4HfNWIPYnNAf3VF0h02E,8768
+waterfall-0.2.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+waterfall-0.2.11.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.2.11.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.2.11.dist-info/RECORD,,

{waterfall-0.2.9.dist-info → waterfall-0.2.11.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any

waterfall-0.2.9.dist-info/RECORD

@@ -1,12 +0,0 @@
-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=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,,