waterfall 0.1.3__tar.gz → 0.1.5__tar.gz

{waterfall-0.1.3 → waterfall-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.1.3
+Version: 0.1.5
 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.3 → waterfall-0.1.5}/README.md

@@ -121,7 +121,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.3 → waterfall-0.1.5}/pyproject.toml

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

waterfall-0.1.5/test.ipynb

@@ -0,0 +1,129 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "ddca610d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from waterfall.watermark import *\n",
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "f0b0a02a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name_or_path = 'meta-llama/Llama-3.1-8B-Instruct'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "ca5f7d55",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2ef1e58174204e879a462a2a829d74fa",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    model_name_or_path,\n",
+    "    torch_dtype=torch.bfloat16,\n",
+    "    device_map=\"cuda:5\",\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "2c475de5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "watermarker = Watermarker(tokenizer=tokenizer, model=model, id=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "e74a58b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "texts = [\"James is a scary man who loves to scare children\", \"There were 3 birds that ate the whole bowl\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "fd0af01e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "watermarked_text = watermark_texts(texts, id, watermarker=watermarker)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "4d58ec82",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['James is an intimidating individual known for taking great pleasure in frightening kids, and this tendency often sends them running in terror.',\n",
+       " \"A small flock comprised of three birds consumed the entirety of the offered bowl's contents within a single meal session.\"]"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "watermarked_text"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

waterfall-0.1.5/test_.ipynb

@@ -0,0 +1,118 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "id": "116f61a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "from itertools import product, combinations\n",
+    "from scipy.stats import binom"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "id": "53c66973",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "np.float64(0.38742048899999965)"
+      ]
+     },
+     "execution_count": 50,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "binom.pmf(1, 10, 0.1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 56,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 32/32 [00:00<00:00, 10845.88it/s]\n",
+      "100%|██████████| 32/32 [00:00<00:00, 3035.09it/s]\n",
+      "100%|██████████| 32/32 [00:00<00:00, 1622.42it/s]\n",
+      "100%|██████████| 32/32 [00:00<00:00, 2020.93it/s]\n",
+      "100%|██████████| 32/32 [00:00<00:00, 4749.72it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0.6582226332676197\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "n = 5   # n-gram\n",
+    "# k = 1   # no of bits modified\n",
+    "l = 8   # bits of watermark\n",
+    "p = 0.1 # probability of modification\n",
+    "\n",
+    "result = 0\n",
+    "for k in range(0, n + 1):\n",
+    "    binomial_prob = binom.pmf(k, n, p)\n",
+    "    if k == 0:\n",
+    "        result += binomial_prob * 1\n",
+    "        continue\n",
+    "    results = []\n",
+    "\n",
+    "    original = np.zeros(n, dtype=bool)\n",
+    "    for original in tqdm(product([True, False], repeat=n), total=2**n):\n",
+    "        original = np.array(original, dtype=bool)\n",
+    "        original_agreement = original.mean() > 0.5\n",
+    "        for modified_indices in combinations(range(n), k):\n",
+    "            modified = original.copy()\n",
+    "            for i in product([True, False], repeat=k):\n",
+    "                modified[np.array(modified_indices)] = i\n",
+    "                modified_agreement = modified.mean() > 0.5\n",
+    "                results.append(original_agreement == modified_agreement)\n",
+    "    result += binomial_prob * np.mean(results) ** l\n",
+    "print(result)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

{waterfall-0.1.3 → waterfall-0.1.5}/waterfall/WatermarkerBase.py

@@ -141,6 +141,7 @@ class Watermarker:
                 do_sample=do_sample,
                 logits_processor=logits_processor,
                 pad_token_id=self.tokenizer.eos_token_id,
+                tokenizer=self.tokenizer,
                 **kwargs
                 )
         output = output[:,tokd_input["input_ids"].shape[-1]:].cpu()

{waterfall-0.1.3 → waterfall-0.1.5}/waterfall/watermark.py

@@ -1,10 +1,12 @@
 import argparse
 import logging
 import os
+import gc
 import torch
 from typing import List, Literal, Optional, Tuple
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers.modeling_utils import PreTrainedModel
 from sentence_transformers import SentenceTransformer
 from tqdm.auto import tqdm
 
@@ -21,6 +23,8 @@ PROMPT = (
 )
 PRE_PARAPHRASED = "Here is a paraphrased version of the text while preserving the semantic similarity:\n\n"
 
+waterfall_cached_watermarking_model = None  # Global variable to cache the watermarking model
+
 def detect_gpu() -> str:
     """
     Use torch to detect if MPS, CUDA, or neither (default CPU)
@@ -42,9 +46,10 @@ def watermark(
     sts_model: SentenceTransformer,
     num_beam_groups: int = 4,
     beams_per_group: int = 2,
-    STS_scale:float = 2.0,
+    STS_scale: float = 2.0,
     diversity_penalty: float = 0.5,
     max_new_tokens: Optional[int] = None,
+    **kwargs
 ) -> str:
     paraphrasing_prompt = watermarker.tokenizer.apply_chat_template(
         [
@@ -61,6 +66,7 @@ def watermark(
         num_beam_groups = num_beam_groups,
         num_return_sequences = num_beam_groups * beams_per_group,
         diversity_penalty = diversity_penalty,
+        **kwargs,
         )
 
     # Select best paraphrasing based on q_score and semantic similarity
@@ -75,7 +81,8 @@ 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"
+                     model_path: Optional[str] = "meta-llama/Llama-3.1-8B-Instruct",
+                     return_extracted_k_p: bool = False
                      ) -> Tuple[float,float]:
     """Returns the q_score and extracted k_p"""
 
@@ -87,7 +94,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"]
 
@@ -140,6 +151,7 @@ def watermark_texts(
     diversity_penalty: float = 0.5,
     STS_scale:float = 2.0,
     use_tqdm: bool = False,
+    stop_at_double_newline: bool = True,    # if True, will stop generation at the first double newline. Prevent repeated paraphrasing of the same text.
 ) -> List[str]:
     if watermark_fn == 'fourier':
         watermarkingFnClass = WatermarkingFnFourier
@@ -150,11 +162,24 @@ def watermark_texts(
 
     if watermarker is None:
         assert model_path is not None, "model_path must be provided if watermarker is not passed"
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            torch_dtype=torch_dtype,
-            device_map=device,
-            )
+        global waterfall_cached_watermarking_model
+
+        if isinstance(waterfall_cached_watermarking_model, PreTrainedModel) and waterfall_cached_watermarking_model.name_or_path != model_path:
+            device = waterfall_cached_watermarking_model.device.type
+            waterfall_cached_watermarking_model = None
+            gc.collect()
+            if device == "cuda":
+                torch.cuda.empty_cache()
+            elif device == "mps":
+                torch.mps.empty_cache()
+
+        if waterfall_cached_watermarking_model is None:
+            waterfall_cached_watermarking_model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=torch_dtype,
+                device_map=device,
+                )
+        model = waterfall_cached_watermarking_model
         tokenizer = AutoTokenizer.from_pretrained(model_path)
 
         watermarker = Watermarker(tokenizer=tokenizer, model=model, id=id, kappa=kappa, k_p=k_p, watermarkingFnClass=watermarkingFnClass)
@@ -173,6 +198,9 @@ def watermark_texts(
     T_ws = []
 
     for T_o in tqdm(T_os, desc="Watermarking texts",  disable=not use_tqdm):
+        if stop_at_double_newline and "\n\n" in T_o:
+            logging.warning("Text contains \\n\\n and stop_at_double_newline is set to True, replacing all \\n\\n in text.")
+            T_o = T_o.replace("\n\n", " ")  # replace double newlines with space
         T_w = watermark(
             T_o,
             watermarker = watermarker,
@@ -181,6 +209,7 @@ def watermark_texts(
             beams_per_group = beams_per_group,
             diversity_penalty = diversity_penalty,
             STS_scale = STS_scale,
+            stop_strings=["\n\n"] if stop_at_double_newline else None,
             )
         T_ws.append(T_w)
 
@@ -283,7 +312,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