waterfall 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

waterfall/watermark.py

@@ -3,6 +3,7 @@ import logging
 import os
 import gc
 import torch
+import numpy as np
 from typing import List, Literal, Optional, Tuple
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -81,8 +82,9 @@ def watermark(
 def verify_texts(texts: List[str], id: int, 
                      watermarker: Optional[Watermarker] = None, 
                      k_p: Optional[int] = None, 
-                     model_path: Optional[str] = "meta-llama/Llama-3.1-8B-Instruct"
-                     ) -> Tuple[float,float]:
+                     model_path: Optional[str] = "meta-llama/Llama-3.1-8B-Instruct",
+                     return_extracted_k_p: bool = False
+                     ) -> np.ndarray | Tuple[np.ndarray,np.ndarray]:
     """Returns the q_score and extracted k_p"""
 
     if watermarker is None:
@@ -93,7 +95,11 @@ def verify_texts(texts: List[str], id: int,
     if k_p is None:
         k_p = watermarker.k_p
 
-    verify_results = watermarker.verify(texts, id=[id], k_p=[k_p], return_extracted_k_p=True)  # results are [text x id x k_p]
+    verify_results = watermarker.verify(texts, id=[id], k_p=[k_p], return_extracted_k_p=return_extracted_k_p)  # results are [text x id x k_p]
+
+    if not return_extracted_k_p:
+        return verify_results[:,0,0]
+
     q_score = verify_results["q_score"]
     k_p_extracted = verify_results["k_p_extracted"]
 
@@ -161,13 +167,12 @@ def watermark_texts(
 
         if isinstance(waterfall_cached_watermarking_model, PreTrainedModel) and waterfall_cached_watermarking_model.name_or_path != model_path:
             device = waterfall_cached_watermarking_model.device.type
-            del waterfall_cached_watermarking_model
+            waterfall_cached_watermarking_model = None
             gc.collect()
             if device == "cuda":
                 torch.cuda.empty_cache()
             elif device == "mps":
                 torch.mps.empty_cache()
-            waterfall_cached_watermarking_model = None
 
         if waterfall_cached_watermarking_model is None:
             waterfall_cached_watermarking_model = AutoModelForCausalLM.from_pretrained(
@@ -308,7 +313,7 @@ def main():
         )
 
     # watermarker = Watermarker(tokenizer=tokenizer, model=None, id=id, k_p=k_p, watermarkingFnClass=watermarkingFnClass)   # If only verifying the watermark, do not need to instantiate the model
-    q_scores, extracted_k_ps = verify_texts(T_os + T_ws, id, watermarker, k_p=k_p)
+    q_scores, extracted_k_ps = verify_texts(T_os + T_ws, id, watermarker, k_p=k_p, return_extracted_k_p=True)
 
     for i in range(len(T_os)):
         # Handle the case where this is being run

{waterfall-0.1.4.dist-info → waterfall-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.1.4
+Version: 0.1.6
 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
@@ -141,7 +141,7 @@ from waterfall.watermark import verify_texts
 id = 1                        # specify your watermarking ID
 test_texts = ["...", "..."]   # Suspected texts to verify
 
-watermark_strength = verify_texts(test_texts, id)[0]   # np array of floats
+watermark_strength = verify_texts(test_texts, id)     # np array of floats
 ```
 
 ## Code structure

{waterfall-0.1.4.dist-info → waterfall-0.1.6.dist-info}/RECORD

@@ -4,9 +4,9 @@ waterfall/WatermarkingFnFourier.py,sha256=QYayAQYwi1dQkDIyqmvhU568VhrVYTVy47HkI8
 waterfall/WatermarkingFnSquare.py,sha256=2PAO05DdKT02npo7GDf_82D520nP7kGAWK6H4E4JMt4,1638
 waterfall/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 waterfall/permute.py,sha256=RwxOHFhx_VSOhhFwy5s79YgwTUBkfW2-LCCXYR3VT2o,2582
-waterfall/watermark.py,sha256=h7e1z8vWTUAKxCcQsJ2Jkx_1ZL-ug2dEDs5FzWcYfCs,13332
-waterfall-0.1.4.dist-info/METADATA,sha256=3hBQwb1JyrTWrayLCPFxXVlTpPjuE-ukPstW5F9F9rg,8715
-waterfall-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-waterfall-0.1.4.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
-waterfall-0.1.4.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
-waterfall-0.1.4.dist-info/RECORD,,
+waterfall/watermark.py,sha256=W5jYGqYGOXXO-KLPKzJoin5zC_Xb6Xk9BzsAA9-LKXA,13494
+waterfall-0.1.6.dist-info/METADATA,sha256=inB0EWGyaAkNQo4CUSke5Bh4Vr07AuVPU9jmKB0gU3U,8714
+waterfall-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+waterfall-0.1.6.dist-info/entry_points.txt,sha256=XXnUzuWXu2nc9j4WAll9tq6HyodN_8WJLjeG0O4Y2Gw,60
+waterfall-0.1.6.dist-info/licenses/LICENSE,sha256=zAtaO-k41Q-Q4Etl4bzuh7pgNJsPH-dYfzvznRa0OvM,11341
+waterfall-0.1.6.dist-info/RECORD,,