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

waterfall/WatermarkerBase.py

@@ -24,6 +24,17 @@ from waterfall.WatermarkingFnFourier import WatermarkingFnFourier
 
 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")
+# Set model loading kwargs based on transformers version
+if version.parse(transformers.__version__) >= version.parse("4.56.0"):
+    model_from_pretrained_kwargs = {"dtype": "auto"}
+else:
+    model_from_pretrained_kwargs = {"torch_dtype": torch.bfloat16}
+
 class PerturbationProcessor(LogitsProcessor):
     def __init__(self,
                  N : int = 32000,     # Vocab size
@@ -134,7 +145,7 @@ class Watermarker:
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name_or_path,
             device_map=device_map,
-            torch_dtype=dtype,
+            **model_from_pretrained_kwargs,
         )
 
     def compute_phi(self, watermarkingFnClass = WatermarkingFnFourier) -> None:
@@ -279,6 +290,9 @@ class Watermarker:
             diversity_penalty = kwargs.get("diversity_penalty", None)
             if num_beams <= 1:
                 kwargs["diversity_penalty"] = None
+        if use_custom_group_beam_search:
+            kwargs["custom_generate"]="transformers-community/group-beam-search"
+            kwargs["trust_remote_code"]=True
 
         if num_beams > 1 and temperature is not None and temperature != 1.0:
             logits_processor.append(TemperatureLogitsWarper(float(temperature)))

waterfall/watermark.py

@@ -15,6 +15,14 @@ from waterfall.WatermarkingFnFourier import WatermarkingFnFourier
 from waterfall.WatermarkingFnSquare import WatermarkingFnSquare
 from waterfall.WatermarkerBase import Watermarker
 
+# Check transformers version
+import transformers
+from packaging import version
+if version.parse(transformers.__version__) >= version.parse("4.56.0"):
+    model_from_pretrained_kwargs = {"dtype": "auto"}
+else:
+    model_from_pretrained_kwargs = {"torch_dtype": torch.bfloat16}
+
 PROMPT = (
     "Paraphrase the user provided text while preserving semantic similarity. "
     "Do not include any other sentences in the response, such as explanations of the paraphrasing. "
@@ -306,8 +314,8 @@ def main():
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     model = AutoModelForCausalLM.from_pretrained(
         model_name_or_path,
-        torch_dtype=torch.bfloat16,
         device_map=device,
+        **model_from_pretrained_kwargs,
         )
 
     watermarker = Watermarker(tokenizer=tokenizer, model=model, id=id, kappa=kappa, k_p=k_p, watermarkingFnClass=watermarkingFnClass)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.2.11
+Version: 0.2.12
 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.57.0,>=4.43.1
+Requires-Dist: transformers>=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 → waterfall-0.2.12.dist-info}/RECORD

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=pk2NB7J0oBLXcO0FIRBHllnSowbpeRNd9ZjvPuUOeeM,21945
+waterfall/WatermarkerBase.py,sha256=6O_S78dD3Jha2OkJK2u3euwCH93i-mTiYYGXosPDMig,22632
 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/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,,