waterfall 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

waterfall/WatermarkerBase.py

@@ -92,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
@@ -116,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,
+            torch_dtype=dtype,
         )
 
     def compute_phi(self, watermarkingFnClass = WatermarkingFnFourier) -> None:

{waterfall-0.2.1.dist-info → waterfall-0.2.3.dist-info}/METADATA

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

@@ -1,12 +1,12 @@
-waterfall/WatermarkerBase.py,sha256=A2VRfsnBfz6-8DSL2NKQZdM1OLI0sQ73qjYaV6rIgJ0,20822
+waterfall/WatermarkerBase.py,sha256=nx2HhNCNy3yAhttq77cFnMmxlGucV_T8iVfbviAmClI,21046
 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.1.dist-info/METADATA,sha256=Mzyp7Nw395RLCN3wnzp2StEpKZEN2erb5BvCOd5Z-4I,8722
-waterfall-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-waterfall-0.2.1.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.2.1.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.2.1.dist-info/RECORD,,
+waterfall-0.2.3.dist-info/METADATA,sha256=8FvklSdyfGq495wvxHZeTT-8nb2xF0tzsH64_tbv6xc,8722
+waterfall-0.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.2.3.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.2.3.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.2.3.dist-info/RECORD,,