waterfall 0.2.6__tar.gz → 0.2.8__tar.gz

{waterfall-0.2.6 → waterfall-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.2.6
+Version: 0.2.8
 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 → waterfall-0.2.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "waterfall"
-version = "0.2.6"
+version = "0.2.8"
 authors = [
   { name = "Xinyuan Niu", email="aperture@outlook.sg" }
 ]

{waterfall-0.2.6 → waterfall-0.2.8}/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] = 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()