waterfall 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

waterfall/WatermarkerBase.py

@@ -22,6 +22,8 @@ from waterfall.permute import Permute
 from waterfall.WatermarkingFn import WatermarkingFn
 from waterfall.WatermarkingFnFourier import WatermarkingFnFourier
 
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 class PerturbationProcessor(LogitsProcessor):
     def __init__(self,
                  N : int = 32000,     # Vocab size
@@ -90,10 +92,9 @@ class Watermarker:
                  n_gram : int = 2,
                  watermarkingFnClass = WatermarkingFnFourier,
                  device = None,
+                 dtype = torch.bfloat16,
                  ) -> None:
         assert kappa >= 0, f"kappa must be >= 0, value provided is {kappa}"
-
-        self.id = id
         self.k_p = k_p
         self.n_gram = n_gram
         self.kappa = kappa
@@ -114,21 +115,26 @@ class Watermarker:
 
         self.N = self.tokenizer.vocab_size
 
-        self.logits_processor = PerturbationProcessor(N = self.N, id = self.id)
-
         if isinstance(model, str):
-            self.load_model(model, device_map=device)
+            self.load_model(model, device_map=device, dtype=dtype)
         else:
             self.model = model
 
         assert (self.model is None) or isinstance(self.model, PreTrainedModel), f"model must be a transformers model, value provided is {type(self.model)}" # argument order for tokenizer and model were swapped since the original code
 
-        self.compute_phi(watermarkingFnClass)
+        self.watermarkingFnClass = watermarkingFnClass
+        self.set_id(id)
+
+    def set_id(self, id : int):
+        self.id = id
+        self.logits_processor = PerturbationProcessor(N = self.N, id = self.id)
+        self.compute_phi(self.watermarkingFnClass)
 
-    def load_model(self, model_name_or_path : str, device_map : str = "auto"):
+    def load_model(self, model_name_or_path : str, device_map : str = "auto", dtype = torch.bfloat16):
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name_or_path,
             device_map=device_map,
+            dtype=dtype,
         )
 
     def compute_phi(self, watermarkingFnClass = WatermarkingFnFourier) -> None:

waterfall/watermark.py

@@ -15,8 +15,6 @@ from waterfall.WatermarkingFnFourier import WatermarkingFnFourier
 from waterfall.WatermarkingFnSquare import WatermarkingFnSquare
 from waterfall.WatermarkerBase import Watermarker
 
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
 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. "

{waterfall-0.2.0.dist-info → waterfall-0.2.2.dist-info}/METADATA

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

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=Jg9jnSU_JYfyCJAyAUqSBcE_IoAByneLC1yKfPEv4mo,20774
+waterfall/WatermarkerBase.py,sha256=cJN1gKY8aFrFf3YVg4H6ukxnc4qLPAtFKyjkSsJjsdQ,21040
 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=_foE_9K1xBdJhTZq2EOeuanl8X04A3wh_F02_1m0LMA,14468
-waterfall-0.2.0.dist-info/METADATA,sha256=AFjq_Ox5oTq2V2a7q80CEAGlYtjpYfOtvEQTSp98w-A,8722
-waterfall-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-waterfall-0.2.0.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.2.0.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.2.0.dist-info/RECORD,,
+waterfall/watermark.py,sha256=IbH5r3oqjtKztDVryfDTr_NDn-CLZHow0S8nAEtZmdc,14420
+waterfall-0.2.2.dist-info/METADATA,sha256=Np7WY47ZGHDGsquIQIxEwWeoOMuDjPuIE8ZBlNCvUQI,8722
+waterfall-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.2.2.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.2.2.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.2.2.dist-info/RECORD,,