waterfall 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

waterfall/WatermarkerBase.py

@@ -227,6 +227,7 @@ class Watermarker:
             return_scores : bool = False,
             use_tqdm : bool = False,
             batched_generate : bool = True,
+            discard_incomplete : bool = True,
             **kwargs    # Other generate parameters
             ) -> List[str] | dict:  # Returns flattened list of query x beam
 
@@ -305,7 +306,10 @@ class Watermarker:
                     tokenizer=self.tokenizer,
                     **kwargs
                     )
-                output = output[:,tokd_input_batch["input_ids"].shape[-1]:].to("cpu", non_blocking=True)
+                output = output[:,tokd_input_batch["input_ids"].shape[-1]:]
+                if discard_incomplete:
+                    output[output[:, -1] == self.tokenizer.eos_token_id].fill_(self.tokenizer.eos_token_id)
+                output = output.to("cpu", non_blocking=True)
                 outputs.append(output)
                 bar.update(tokd_input_batch["input_ids"].shape[0])
         torch.cuda.synchronize()

{waterfall-0.2.6.dist-info → waterfall-0.2.7.dist-info}/METADATA

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

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=WAngzoYKOZ9FMCP9jZGLChxpPMLYum05EsEbRsN59Bc,21080
+waterfall/WatermarkerBase.py,sha256=_Ok-lU-gF3xEWCBwEmxW6ncvUaFXdGQbdiBlAqFAhno,21305
 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.6.dist-info/METADATA,sha256=ealMvDHrvzovcCpIsqjpBxZ3ukbN0WXDxdgrQYkgnDE,8723
-waterfall-0.2.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-waterfall-0.2.6.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.2.6.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.2.6.dist-info/RECORD,,
+waterfall-0.2.7.dist-info/METADATA,sha256=GhqFgFaxYJD11tKHejrS9PsHBL2cdggieezeXWvLdxk,8723
+waterfall-0.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.2.7.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.2.7.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.2.7.dist-info/RECORD,,