waterfall 0.2.5__tar.gz → 0.2.7__tar.gz

{waterfall-0.2.5 → waterfall-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: waterfall
-Version: 0.2.5
+Version: 0.2.7
 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.5 → waterfall-0.2.7}/pyproject.toml

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

{waterfall-0.2.5 → waterfall-0.2.7}/waterfall/WatermarkerBase.py

@@ -227,6 +227,7 @@ class Watermarker:
             return_scores : bool = False,
             use_tqdm : bool = False,
             batched_generate : bool = True,
+            discard_incomplete : bool = True,
             **kwargs    # Other generate parameters
             ) -> List[str] | dict:  # Returns flattened list of query x beam
 
@@ -305,7 +306,10 @@ class Watermarker:
                     tokenizer=self.tokenizer,
                     **kwargs
                     )
-                output = output[:,tokd_input_batch["input_ids"].shape[-1]:].to("cpu", non_blocking=True)
+                output = output[:,tokd_input_batch["input_ids"].shape[-1]:]
+                if discard_incomplete:
+                    output[output[:, -1] == self.tokenizer.eos_token_id].fill_(self.tokenizer.eos_token_id)
+                output = output.to("cpu", non_blocking=True)
                 outputs.append(output)
                 bar.update(tokd_input_batch["input_ids"].shape[0])
         torch.cuda.synchronize()

waterfall-0.2.5/.vscode/settings.json

@@ -1,5 +0,0 @@
-{
-    "python-envs.defaultEnvManager": "ms-python.python:conda",
-    "python-envs.defaultPackageManager": "ms-python.python:conda",
-    "python-envs.pythonProjects": []
-}

waterfall-0.2.5/pypi.txt

@@ -1,2 +0,0 @@
-python3 -m build
-python3 -m twine upload dist/*

waterfall-0.2.5/test.ipynb

@@ -1,129 +0,0 @@
-{
- "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.2.5/test_.ipynb

@@ -1,118 +0,0 @@
-{
- "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
-}