From 09e64c4dedb3f752950c0a423da852cb4057b8e3 Mon Sep 17 00:00:00 2001 From: "Guillaume Meyer (The Opinionated Man)" <1385518+guillaumemeyer@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:56:16 -0700 Subject: [PATCH] feat: optional MarkLLM text-watermark verification harness (#53) * feat: optional MarkLLM text-watermark verification harness Add an optional external backend wrapping THU-BPM/MarkLLM (Apache-2.0) so a specific statistical text-watermark scheme (KGW / SynthID-Text) can be verified before/after a Layer B rewrite. - detect_text_watermark.py: detect/watermark subcommands, external checkout at a pinned commit, exit codes 0/1/2/3, --json - rewrite_text.py --markllm-scheme: before/after detection around the rewrite, reports a `cleared` flag; never fails the rewrite when the backend is unavailable - setup_markllm.sh + requirements-markllm.txt (pinned deps) + Dockerfile.markllm + Makefile bootstrap/smoke/docker targets - tests/test_markllm_detect.py: 16 mock-based cases (no torch in CI) - Docs: verification-harness caveat (same-config-only, not a vendor-detector oracle) in README, SKILL.md, removal-matrix, vendor-notes * chore: tidy merged Unreleased changelog list * security: harden the MarkLLM harness (offline, caps, supply-chain) Addresses the PR security review: - detect_text_watermark.py: --offline loads the scoring model from the HF cache only (local_files_only + HF_HUB_OFFLINE, no remote code), and the algorithm config is capped at 1 MiB so a crafted huge file is refused before either this script or upstream reads it into memory - rewrite_text.py: WATERMARKS_MARKLLM_RLIMIT_AS (env, POSIX) optionally applies RLIMIT_AS to the MarkLLM subprocess; off by default because torch/CUDA needs large address spaces - Dockerfile.markllm: drop the unpinned torch install (it is pinned in requirements-markllm.txt) and verify the cloned upstream commit SHA - tests: offline flag, config-too-large, and preexec/rlimit cases - docs: hardening knobs in README + SKILL.md; changelog updated --- Dockerfile.markllm | 69 +++ Makefile | 19 +- README.md | 72 +++ skills/remove-ai-marks/SKILL.md | 23 +- .../references/removal-matrix.md | 3 +- .../references/vendor-notes.md | 2 + .../scripts/detect_text_watermark.py | 345 ++++++++++++++ .../scripts/requirements-markllm.txt | 27 ++ .../remove-ai-marks/scripts/rewrite_text.py | 166 +++++++ .../remove-ai-marks/scripts/setup_markllm.sh | 132 ++++++ tests/test_markllm_detect.py | 435 ++++++++++++++++++ 11 files changed, 1290 insertions(+), 3 deletions(-) create mode 100644 Dockerfile.markllm create mode 100755 skills/remove-ai-marks/scripts/detect_text_watermark.py create mode 100644 skills/remove-ai-marks/scripts/requirements-markllm.txt create mode 100755 skills/remove-ai-marks/scripts/setup_markllm.sh create mode 100644 tests/test_markllm_detect.py diff --git a/Dockerfile.markllm b/Dockerfile.markllm new file mode 100644 index 0000000..d854358 --- /dev/null +++ b/Dockerfile.markllm @@ -0,0 +1,69 @@ +# Optional local Docker image for the MarkLLM text-watermark harness. +# +# Build from the repository root: +# docker build -f Dockerfile.markllm -t watermarks-remover-markllm . +# +# The upstream code is fetched from source at build time and is NOT +# redistributed by this repository. Upstream is Apache-2.0. +# +# Vendored fork hardening: +# - base image pinned by digest (no moving tag drift) +# - upstream checkout pinned to a commit SHA (no moving branch) +# - deps pinned exactly in requirements-markllm.txt +# - pip itself pinned (no unpinned bootstrap step) +# - runs as an unprivileged user (a parser bug in a crafted file can no +# longer write files as root inside the container) + +# Pinned upstream commit (2026-07-10). Keep in sync with setup_markllm.sh. +ARG MARKLLM_REF=c45ddc40f7b761beabe55a1b8dc4690e531d1c6d + +# python:3.14-slim linux/amd64 digest. +FROM python:3.14-slim@sha256:ce40764625a4ff50df3548277632e7f96c4e77fe75fa848aae9885476e7df5a4 + +ARG MARKLLM_REF + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + git \ + libgl1 \ + libglib2.0-0 \ + passwd \ + && rm -rf /var/lib/apt/lists/* + +RUN git clone --depth 1 --filter=blob:none --sparse \ + https://github.com/THU-BPM/MarkLLM.git /opt/markllm \ + && cd /opt/markllm \ + && git fetch --depth 1 origin "${MARKLLM_REF}" \ + && git checkout --detach "${MARKLLM_REF}" \ + && git sparse-checkout set --no-cone \ + '/watermark/' \ + '/config/' \ + '/utils/' \ + '/exceptions/' \ + '/evaluation/dataset.py' \ + '/LICENSE' \ + '/README.md' \ + && test "$(git -C /opt/markllm rev-parse HEAD)" = "${MARKLLM_REF}" + +COPY skills/remove-ai-marks/scripts/requirements-markllm.txt /app/requirements-markllm.txt +COPY skills/remove-ai-marks/scripts/detect_text_watermark.py /app/detect_text_watermark.py + +# torch is pinned (with everything else) in requirements-markllm.txt; there is +# no separate unpinned install step. No GPU wheel index inside the image — +# CUDA users should run setup_markllm.sh on the host instead. +RUN python3 -m pip install --no-cache-dir "pip==26.2.1" \ + && python3 -m pip install --no-cache-dir -r /app/requirements-markllm.txt + +# Unprivileged runtime user. The harness only reads input files and writes to +# stdout, so nothing under /opt, /app, or the mounted data dir needs root. +RUN useradd --create-home --uid 10001 --shell /usr/sbin/nologin markllm +USER markllm + +ENV MARKLLM_DIR=/opt/markllm \ + HOME=/home/markllm \ + PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + HF_HOME=/home/markllm/.cache/huggingface + +WORKDIR /app +ENTRYPOINT ["python3", "/app/detect_text_watermark.py"] diff --git a/Makefile b/Makefile index 561e23a..274333a 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,6 @@ .PHONY: test smoke smoke-synthid bootstrap-synthid docker-synthid-build docker-synthid-help \ - smoke-ctrlregen bootstrap-ctrlregen docker-ctrlregen-build docker-ctrlregen-help install-skill clean + smoke-ctrlregen bootstrap-ctrlregen docker-ctrlregen-build docker-ctrlregen-help \ + smoke-markllm bootstrap-markllm docker-markllm-build docker-markllm-help install-skill clean SCRIPTS := skills/remove-ai-marks/scripts PYTHON ?= $(shell if [ -x .venv/bin/python ]; then echo .venv/bin/python; else echo python3; fi) @@ -49,6 +50,22 @@ docker-ctrlregen-build: docker-ctrlregen-help: docker run --rm watermarks-remover-ctrlregen --help +smoke-markllm: + @if [ -z "$(MARKLLM_DIR)" ]; then \ + echo "smoke-markllm skipped (set MARKLLM_DIR)"; \ + else \ + $(PYTHON) $(SCRIPTS)/detect_text_watermark.py --help >/dev/null && echo "detect_text_watermark adapter present"; \ + fi + +bootstrap-markllm: + ./skills/remove-ai-marks/scripts/setup_markllm.sh + +docker-markllm-build: + docker build -f Dockerfile.markllm -t watermarks-remover-markllm . + +docker-markllm-help: + docker run --rm watermarks-remover-markllm --help + install-skill: mkdir -p $(HOME)/.grok/skills ln -sfn $(CURDIR)/skills/remove-ai-marks $(HOME)/.grok/skills/remove-ai-marks diff --git a/README.md b/README.md index 2b6b8f1..d493be2 100644 --- a/README.md +++ b/README.md @@ -241,6 +241,72 @@ docker run --rm -e HF_TOKEN="$HF_TOKEN" \ watermarks-remover-ctrlregen /data/shot.png -o /data/shot.ctrlregen.png ``` +## Optional MarkLLM text-watermark verification + +For **controlled experiments**, an optional external harness wraps +[`THU-BPM/MarkLLM`](https://github.com/THU-BPM/MarkLLM) (Apache-2.0) to +watermark test text and re-detect it after a Layer B rewrite — e.g. prove that +a KGW (Kirchenbauer, your "open-LLM" row) or SynthID-Text (Gemini row) mark +disappears under your rewrite. It is a **verification harness, not an oracle**: +MarkLLM detection is only valid against the *same* scheme config + keys used at +generation, and it cannot certify a vendor detector will fail. + +The backend is **not bundled**. `setup_markllm.sh` clones upstream at a pinned +commit, creates a venv, and installs pinned deps (torch + transformers); the +scoring model (default `facebook/opt-1.3b`, Apache-2.0) downloads from Hugging +Face on first run. + +```bash +SCRIPTS=skills/remove-ai-marks/scripts + +# Bootstrap (clones upstream, creates ~/MarkLLM/.venv, installs deps). +"$SCRIPTS/setup_markllm.sh" + +# Generate watermarked + unwatermarked sample text under the KGW scheme. +MARKLLM_DIR=~/MarkLLM \ + ~/MarkLLM/.venv/bin/python "$SCRIPTS/detect_text_watermark.py" watermark prompt.txt \ + --scheme kgw -o wm.txt -o2 plain.txt + +# Detect the scheme mark in a text file. +MARKLLM_DIR=~/MarkLLM \ + ~/MarkLLM/.venv/bin/python "$SCRIPTS/detect_text_watermark.py" detect wm.txt --scheme kgw --json +``` + +**Verification around a Layer B rewrite:** pass `--markllm-scheme` to +`rewrite_text.py` (with `--markllm-dir`), and it records the MarkLLM detection +before/after plus a `cleared` flag: + +```bash +export WATERMARKS_REWRITE_BACKEND=ollama WATERMARKS_REWRITE_MODEL=llama3.2 +MARKLLM_DIR=~/MarkLLM \ + python3 "$SCRIPTS/rewrite_text.py" wm.txt -o wm.rewritten.txt \ + --markllm-scheme kgw --markllm-dir "$HOME/MarkLLM" --json-stats +``` + +If the backend is unconfigured or its deps are missing, the rewrite proceeds +and the report notes verification was unavailable. A GPU is recommended; CPU +runs work but are slow, and the model download is a few GB. + +Hardening knobs: + +- `--offline` on the adapter (or any MarkLLM run) loads the scoring model from + the Hugging Face cache only — zero network egress; fails fast if not cached. + Custom remote code is never executed (transformers `trust_remote_code` is + never enabled). +- `WATERMARKS_MARKLLM_RLIMIT_AS=` (env, POSIX) applies an address-space + limit to the MarkLLM subprocess spawned by `rewrite_text.py`. Off by default + because torch/CUDA usually needs large address spaces. +- Config files are capped at 1 MiB; the upstream checkout and the base image + are pinned by SHA/digest. + +### Docker + +```bash +make docker-markllm-build +docker run --rm --user "$(id -u):$(id -g)" -v "$(pwd):/data" \ + watermarks-remover-markllm detect /data/wm.txt --scheme kgw --json +``` + ## Coverage matrix | Channel | Claude | Gemini/SynthID | OpenAI | Open-LLM | @@ -369,6 +435,12 @@ make smoke # quick CLI smoke on fixtures ### Unreleased - Add stdlib-only WebP inspection and metadata cleaning for RIFF `C2PA`, XMP, EXIF, and ICC profile chunks +- New optional MarkLLM harness (external `THU-BPM/MarkLLM` checkout, Apache-2.0): `detect_text_watermark.py` with `detect` / `watermark` subcommands for KGW and SynthID schemes +- `rewrite_text.py --markllm-scheme` runs before/after detection around a Layer B rewrite (env-gated; reports `cleared`) +- `setup_markllm.sh` bootstrap + `requirements-markllm.txt` (pinned deps) + `Dockerfile.markllm` and Makefile `bootstrap-markllm` / `smoke-markllm` / `docker-markllm-build` / `docker-markllm-help` +- Mock-based tests (`tests/test_markllm_detect.py`, 21 cases) — no torch in CI +- Docs: verification-harness caveat (same-config-only, not a vendor-detector oracle) in README, SKILL.md, `removal-matrix.md`, `vendor-notes.md` +- Hardening: `--offline` cache-only model loading (no HF egress, no remote code), 1 MiB config cap, optional `WATERMARKS_MARKLLM_RLIMIT_AS` on the rewrite subprocess, pinned torch in the Dockerfile, and clone-SHA verification in `Dockerfile.markllm` ### [v0.4.0](https://github.com/guillaumemeyer/watermarks-remover/releases/tag/v0.4.0) — pixel removal, finding confidence, Windows & false-positive fixes diff --git a/skills/remove-ai-marks/SKILL.md b/skills/remove-ai-marks/SKILL.md index 3c2d3ea..6e2607c 100644 --- a/skills/remove-ai-marks/SKILL.md +++ b/skills/remove-ai-marks/SKILL.md @@ -35,6 +35,8 @@ python3 "$SCRIPTS/clean_image.py" ... python3 "$SCRIPTS/clean_ctrlregen.py" ... # optional external pixel removal (bootstrap first) "$SCRIPTS/setup_ctrlregen.sh" # one-command bootstrap (Windows: setup_ctrlregen.ps1) python3 "$SCRIPTS/rewrite_text.py" ... +python3 "$SCRIPTS/detect_text_watermark.py" ... # optional external MarkLLM verification harness (bootstrap first) +"$SCRIPTS/setup_markllm.sh" # one-command bootstrap for the above python3 "$SCRIPTS/audit_dir.py" ... python3 "$SCRIPTS/audit_website.py" ... ``` @@ -154,6 +156,17 @@ python3 "$SCRIPTS/rewrite_text.py" draft.md -o draft.rewritten.md --strength par If the hook is not configured, run the prompts below yourself (agent-orchestrated). +**Optional MarkLLM verification:** to test a specific scheme (KGW / SynthID) +under a config you control, bootstrap the external MarkLLM checkout with +`scripts/setup_markllm.sh`, then either run +`detect_text_watermark.py watermark`/`detect` directly or pass +`--markllm-scheme kgw|synthid` to `rewrite_text.py` for a before/after +detection report. This is a **controlled-experiment harness** — detection is +only valid against the same scheme config + keys used at generation and cannot +certify a vendor detector. Add `--offline` to load the scoring model from the +HF cache only (no network); `WATERMARKS_MARKLLM_RLIMIT_AS` (env, POSIX) +optionally caps the subprocess address space. + **Code files:** Prefer formatter (`prettier`, `black`, `gofmt`, …) + Layer A. Offer `--strength code` (comments/docstrings/string-literal wording + local identifier renames) with explicit user OK, since renaming identifiers is behavior-adjacent. #### Rewrite prompts (use as-is) @@ -235,7 +248,7 @@ Always state: ## Limitations - Layer A does **not** remove token-sampling watermarks. -- Layer B cannot be gold-verified without vendor detectors / keys. +- Layer B cannot be gold-verified without vendor detectors / keys. The optional MarkLLM harness (`detect_text_watermark.py` / `rewrite_text.py --markllm-scheme`) verifies a specific scheme config before/after a rewrite, but it is same-config-only and not a vendor-detector oracle; its backend is external (Apache-2.0), never bundled, and pulls torch + a few GB of model weights. - PDF strip is best-effort without `exiftool`, and incomplete without `qpdf`: exiftool alone leaves the freed metadata objects in the byte stream. - Pixel-domain **image** watermarks can be removed optionally via the external CtrlRegen backend (`clean_image.py --remove-pixel ctrlregen`); audio/video watermarks remain out of scope for removal. - The CtrlRegen backend is external, all-rights-reserved (no LICENSE file), never bundled, heavy (~10 GB model downloads), and a regenerating remover — no local detector certifies StegaStamp/Tree-Ring/StableSignature removal. @@ -256,6 +269,7 @@ python3 scripts/clean_file.py deck.docx -o deck.cleaned.docx python3 scripts/inspect_text.py notes.md python3 scripts/clean_text.py notes.md -o notes.cleaned.md --stats python3 scripts/rewrite_text.py notes.md --backend print-prompt --strength paraphrase +python3 scripts/rewrite_text.py notes.md --markllm-scheme kgw --markllm-dir ~/MarkLLM # optional before/after verification # Images only python3 scripts/inspect_image.py shot.png @@ -267,6 +281,13 @@ NOAI_WATERMARK_DIR=~/noai-watermark \ ~/noai-watermark/.venv/bin/python scripts/clean_image.py shot.png \ -o shot.cleaned.png --remove-pixel ctrlregen +# Optional MarkLLM verification (external backend; bootstrap first) +scripts/setup_markllm.sh +MARKLLM_DIR=~/MarkLLM \ + ~/MarkLLM/.venv/bin/python scripts/detect_text_watermark.py detect notes.md --scheme kgw +MARKLLM_DIR=~/MarkLLM \ + ~/MarkLLM/.venv/bin/python scripts/detect_text_watermark.py watermark prompt.txt --scheme synthid -o wm.txt + # Aggregate audits python3 scripts/audit_dir.py ./src --json python3 scripts/audit_website.py --sitemap https://example.com/sitemap.xml --json diff --git a/skills/remove-ai-marks/references/removal-matrix.md b/skills/remove-ai-marks/references/removal-matrix.md index db6a8a4..cce3c5b 100644 --- a/skills/remove-ai-marks/references/removal-matrix.md +++ b/skills/remove-ai-marks/references/removal-matrix.md @@ -3,7 +3,7 @@ | Target | Method | Script / action | Side effects | Verifiable today? | | --- | --- | --- | --- | --- | | Invisible Unicode / exotic spaces / bidi / tags | Strip / normalize | `inspect_text.py`, `clean_text.py`, `clean_file.py` | Minimal | Yes (codepoint report) | -| Statistical text watermark (SynthID-class / Kirchenbauer) | Multi-pass paraphrase / humanize / back-translate / structural | Agent Layer B + optional `rewrite_text.py` | Meaning/style drift | No without vendor key/detector | +| Statistical text watermark (SynthID-class / Kirchenbauer) | Multi-pass paraphrase / humanize / back-translate / structural | Agent Layer B + optional `rewrite_text.py` | Meaning/style drift | No without vendor key/detector; **MarkLLM harness** (`detect_text_watermark.py`) verifies a specific scheme config before/after | | C2PA on PNG/JPEG/WebP | Drop APP11 / PNG `caBX` / RIFF `C2PA` / exiftool | `clean_image.py` | Loses provenance metadata | Yes | | SVG metadata / XMP | Drop ``, xmpmeta | `clean_file.py` | Loses SVG metadata | Yes (re-inspect) | | PDF XMP / info | exiftool `-all=` preferred | `clean_file.py` | Loses PDF metadata; degraded without exiftool | Partial | @@ -24,6 +24,7 @@ 4. Prefer a **non-origin, open-weight** rewrite model when available (avoid re-stamping). 5. Layer A again after rewrite. 6. Report: Layer B is best-effort; residual risk remains. +7. **Optional verification:** `rewrite_text.py --markllm-scheme kgw|synthid` runs a MarkLLM before/after detection (external `detect_text_watermark.py` harness) to show a specific scheme config clears. Same-config-only; not a vendor-detector oracle. ## Code vs prose diff --git a/skills/remove-ai-marks/references/vendor-notes.md b/skills/remove-ai-marks/references/vendor-notes.md index d494a65..2c51820 100644 --- a/skills/remove-ai-marks/references/vendor-notes.md +++ b/skills/remove-ai-marks/references/vendor-notes.md @@ -32,6 +32,7 @@ Source: [How Claude marks AI-generated content](https://support.claude.com/en/ar - **Data-driven / backdoor** (trigger phrases) → **out of scope** - **Generative** (sampling) → Layer B best-effort - Productionized in Gemini-scale systems; open research code exists, but **production keys are not public** — this skill does **not** ship a SynthID detector. +- Optional external verification harness: [`THU-BPM/MarkLLM`](https://github.com/THU-BPM/MarkLLM) (Apache-2.0) reimplements SynthID-Text among other schemes with configurable keys; wired as `detect_text_watermark.py` / `rewrite_text.py --markllm-scheme`. Same-config-only — it verifies a mark you generated under a known config, not Google's production keying. - Current frontier production watermarks are **token-by-token** (streaming constraint); paragraph-level robust methods (SemStamp / PostMark) are not deployed yet, which keeps paraphrase-class attacks effective today. - Optional external reference: [`aloshdenny/reverse-SynthID`](https://github.com/aloshdenny/reverse-SynthID) provides a reverse-engineered pixel-domain scorer. It is **not bundled** here, is best-effort, and is under a non-commercial Research License; it is not the official Google detector. - Optional pixel-domain removal: [`mertizci/noai-watermark`](https://github.com/mertizci/noai-watermark)'s CtrlRegen profile is wired through `clean_image.py --remove-pixel ctrlregen` / `clean_ctrlregen.py`. It is **not bundled** (no LICENSE file → all-rights-reserved), and no local detector certifies the result; the official Google check is the final authority. @@ -50,6 +51,7 @@ Source: [How Claude marks AI-generated content](https://support.claude.com/en/ar - Classic green-list / red-list sampling bias (Kirchenbauer et al.) and variants. - Detectable with the **key** and tokenizer; removal still relies on heavy paraphrase or regeneration. +- Optional external harness: `MarkLLM` (`detect_text_watermark.py --scheme kgw`) reproduces KGW detection under a config you control, for controlled before/after experiments. **Skill mapping:** Layer B multi-pass; prefer rewrite with a **different** model family when possible. diff --git a/skills/remove-ai-marks/scripts/detect_text_watermark.py b/skills/remove-ai-marks/scripts/detect_text_watermark.py new file mode 100755 index 0000000..8684490 --- /dev/null +++ b/skills/remove-ai-marks/scripts/detect_text_watermark.py @@ -0,0 +1,345 @@ +#!/usr/bin/env python3 +"""Optional MarkLLM text-watermark harness backed by an external THU-BPM/MarkLLM checkout. + +This script does NOT vendor upstream code. It imports ``AutoWatermark`` from a +user-provided checkout (https://github.com/THU-BPM/MarkLLM) at runtime, using +that environment's optional dependencies (torch, transformers, datasets, ...). + +MarkLLM is Apache-2.0. It is a research/verification harness: detection is only +valid against the SAME scheme config + keys used at generation. It cannot +certify that a vendor detector will fail on the given text. + +Subcommands: + detect run detection on a text file with a known scheme/config + watermark generate watermarked (and optionally unwatermarked) sample text + from a prompt, for controlled before/after experiments + +Exit codes: + 0 success + 1 runtime error (model load, detection/generation failure) + 2 bad input (missing/unreadable file, binary input, bad args) + 3 unavailable (not configured / missing checkout / missing deps) +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +SCRIPTS_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(SCRIPTS_DIR)) + +from common import eprint, emit_json, read_text_input, safe_write_text # noqa: E402 + +# Scheme name as the user types it -> MarkLLM algorithm name (config/{ALG}.json). +SCHEMES = { + "kgw": "KGW", + "synthid": "SynthID", + "synthid-text": "SynthID", +} + +DEFAULT_MODEL = "facebook/opt-1.3b" + +# Algorithm configs are ~200 B (KGW/SynthID). Cap well above that so a crafted +# or accidental huge file is refused before either this script or upstream +# reads it into memory. +MAX_CONFIG_BYTES = 1 << 20 + + +class _Unavailable(RuntimeError): + """Backend present but unusable (unconfigured checkout, missing deps).""" + + +def resolve_upstream(raw: str | None) -> Path | None: + if not raw: + return None + upstream = Path(raw).expanduser().resolve() + if not upstream.is_dir(): + return None + return upstream + + +def resolve_device(raw: str | None) -> str: + """Resolve the ``auto`` device hint to a concrete torch device.""" + if raw and raw != "auto": + return raw + try: + import torch + + if torch.cuda.is_available(): + return "cuda" + mps = getattr(torch.backends, "mps", None) + if mps is not None and mps.is_available(): + return "mps" + except Exception: + pass + return "cpu" + + +def _load_algorithm( + upstream: Path, alg: str, config: Path, model: str, device: str, offline: bool = False +): + """Import the checkout and build an ``AutoWatermark`` instance.""" + sys.path.insert(0, str(upstream)) + try: + from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 + from utils.transformers_config import TransformersConfig # noqa: E402 + from watermark.auto_watermark import AutoWatermark # noqa: E402 + except ImportError as e: + raise _Unavailable(f"MarkLLM dependencies missing: {e}") from e + + # --offline: never contact the HF hub. local_files_only makes transformers + # fail fast instead of hanging, and HF_HUB_OFFLINE covers the lower-level + # hub calls. Custom-code execution is not possible either way: transformers + # only honors auto_map/trust_remote_code when explicitly enabled, which is + # never done here. + if offline: + os.environ.setdefault("HF_HUB_OFFLINE", "1") + load_kwargs = {"local_files_only": True} if offline else {} + + tokenizer = AutoTokenizer.from_pretrained(model, **load_kwargs) + lm = AutoModelForCausalLM.from_pretrained(model, **load_kwargs).to(device) + transformers_config = TransformersConfig( + model=lm, + tokenizer=tokenizer, + device=device, + max_new_tokens=200, + min_length=0, + do_sample=True, + no_repeat_ngram_size=4, + ) + return AutoWatermark.load( + alg, + algorithm_config=str(config), + transformers_config=transformers_config, + ) + + +def _threshold_from_config(config: Path) -> float | None: + try: + data = json.loads(config.read_text("utf-8")) + except (OSError, ValueError): + return None + for key in ("threshold", "z_threshold"): + value = data.get(key) + if isinstance(value, (int, float)): + return float(value) + return None + + +def _resolve_config(upstream: Path, alg: str, config: str | None) -> Path: + if config: + path = Path(config).expanduser().resolve() + else: + path = upstream / "config" / f"{alg}.json" + if not path.is_file(): + raise _Unavailable(f"MarkLLM config not found: {path}") + try: + size = path.stat().st_size + except OSError as e: + raise _Unavailable(f"cannot stat MarkLLM config {path}: {e}") from e + if size > MAX_CONFIG_BYTES: + raise _Unavailable( + f"MarkLLM config too large ({size} bytes > {MAX_CONFIG_BYTES}): {path}" + ) + return path + + +def _cmd_detect(args: argparse.Namespace, upstream: Path, alg: str) -> int: + if args.path != "-" and not Path(args.path).is_file(): + eprint(f"not a file: {args.path}") + return 2 + text = read_text_input(args.path, allow_binary=args.force_text) + + device = resolve_device(args.device) + + try: + config = _resolve_config(upstream, alg, args.config) + threshold = _threshold_from_config(config) + wm = _load_algorithm( + upstream, alg, config, args.model, device, offline=args.offline + ) + result = wm.detect_watermark(text, return_dict=True) + except _Unavailable as e: + eprint(str(e)) + return 3 + except Exception as e: + eprint(f"detection error: {e}") + return 1 + + is_watermarked = bool(result.get("is_watermarked", False)) + score = result.get("score") + try: + score = float(score) + except (TypeError, ValueError): + score = None + + payload = { + "available": True, + "upstream_dir": str(upstream), + "scheme": alg, + "config": str(config), + "model": args.model, + "device": device, + "is_watermarked": is_watermarked, + "score": score, + "threshold": threshold, + } + + if args.json: + emit_json(payload) + else: + label = "watermarked" if is_watermarked else "not watermarked" + score_txt = f"{score:.4f}" if score is not None else "n/a" + thresh_txt = f"{threshold:.4f}" if threshold is not None else "n/a" + print(f"{alg}: {label} (score {score_txt}, threshold {thresh_txt})") + + return 0 + + +def _cmd_watermark(args: argparse.Namespace, upstream: Path, alg: str) -> int: + prompt = read_text_input(args.prompt, allow_binary=args.force_text) + + device = resolve_device(args.device) + + try: + config = _resolve_config(upstream, alg, args.config) + wm = _load_algorithm( + upstream, alg, config, args.model, device, offline=args.offline + ) + if args.seed is not None: + import torch + + torch.manual_seed(args.seed) + wm.config.gen_kwargs["max_new_tokens"] = args.max_new_tokens + wm.config.gen_kwargs["min_length"] = args.min_length + watermarked = wm.generate_watermarked_text(prompt) + unwatermarked = None + if args.unwatermarked_output: + unwatermarked = wm.generate_unwatermarked_text(prompt) + except _Unavailable as e: + eprint(str(e)) + return 3 + except Exception as e: + eprint(f"generation error: {e}") + return 1 + + wm_out = "-" if args.watermarked_output is None else args.watermarked_output + safe_write_text(wm_out, watermarked) + if unwatermarked is not None: + safe_write_text(args.unwatermarked_output, unwatermarked) + + payload = { + "available": True, + "upstream_dir": str(upstream), + "scheme": alg, + "config": str(config), + "model": args.model, + "device": device, + "watermarked_output": wm_out, + "unwatermarked_output": args.unwatermarked_output, + "watermarked_chars": len(watermarked), + "unwatermarked_chars": len(unwatermarked) if unwatermarked is not None else None, + } + + if args.json: + emit_json(payload) + else: + print(f"{alg}: watermarked sample ({payload['watermarked_chars']} chars) -> {wm_out}") + if unwatermarked is not None: + print(f" unwatermarked sample ({payload['unwatermarked_chars']} chars) -> {args.unwatermarked_output}") + + return 0 + + +def _add_common(p: argparse.ArgumentParser) -> None: + p.add_argument( + "--upstream-dir", + type=Path, + default=None, + help="MarkLLM checkout root (default: $MARKLLM_DIR)", + ) + p.add_argument( + "--scheme", + required=True, + choices=sorted(SCHEMES), + help="Watermark scheme to use (kgw, synthid)", + ) + p.add_argument( + "--config", + default=None, + help="Algorithm config JSON (default: /config/.json)", + ) + p.add_argument( + "--model", + default=os.environ.get("MARKLLM_MODEL", DEFAULT_MODEL), + help=f"HF causal LM for scoring (default: $MARKLLM_MODEL or {DEFAULT_MODEL})", + ) + p.add_argument( + "--device", + default="auto", + help="auto|cpu|cuda|mps (default: auto)", + ) + p.add_argument( + "--offline", + action="store_true", + help="Never contact the HF hub: load the scoring model from the local " + "cache only (fails fast if not cached)", + ) + p.add_argument( + "--force-text", + action="store_true", + help="Process input even when it looks like a binary container", + ) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd", required=True) + + detect = sub.add_parser("detect", help="Detect a scheme watermark in text") + detect.add_argument("path", help="Text file to detect on, or - for stdin") + _add_common(detect) + detect.add_argument("--json", action="store_true", help="Emit JSON on stdout") + detect.set_defaults(handler=_cmd_detect) + + wm = sub.add_parser("watermark", help="Generate watermarked sample text") + wm.add_argument("prompt", help="Prompt file, or - for stdin") + wm.add_argument("-o", "--watermarked-output", default=None, + help="Output path for the watermarked sample (default: stdout)") + wm.add_argument("-o2", "--unwatermarked-output", default=None, + help="Also write an unwatermarked sample to this path") + wm.add_argument("--max-new-tokens", type=int, default=200) + wm.add_argument("--min-length", type=int, default=0) + wm.add_argument("--seed", type=int, default=None, help="Optional RNG seed") + _add_common(wm) + wm.add_argument("--json", action="store_true", help="Emit JSON on stdout") + wm.set_defaults(handler=_cmd_watermark) + + args = p.parse_args() + + if args.cmd == "detect" and args.path != "-" and not Path(args.path).is_file(): + eprint(f"not a file: {args.path}") + return 2 + + raw_upstream = args.upstream_dir or os.environ.get("MARKLLM_DIR") + upstream = resolve_upstream(str(raw_upstream) if raw_upstream else None) + if upstream is None: + eprint( + "MarkLLM not configured: set MARKLLM_DIR or pass --upstream-dir", + ) + return 3 + + if not (upstream / "watermark").is_dir(): + eprint(f"MarkLLM checkout incomplete (no watermark/ dir): {upstream}") + return 3 + + alg = SCHEMES[args.scheme] + return args.handler(args, upstream, alg) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/remove-ai-marks/scripts/requirements-markllm.txt b/skills/remove-ai-marks/scripts/requirements-markllm.txt new file mode 100644 index 0000000..f1fa930 --- /dev/null +++ b/skills/remove-ai-marks/scripts/requirements-markllm.txt @@ -0,0 +1,27 @@ +# Minimal dependencies for the optional MarkLLM text-watermark harness. +# The upstream repo's requirements.txt is unpinned and pulls evaluation-only +# deps (sentence-transformers, sacrebleu, networkx, translate, tiktoken, +# openai==0.28). This file installs only what AutoWatermark.load() needs for +# the KGW and SynthID schemes (plus the synthid detector's C4Dataset import). +# +# Vendored fork hardening: exact pins (no drift). Re-evaluate versions before +# bumping and keep the pinned upstream checkout in setup_markllm.sh / +# Dockerfile.markllm in sync with the code these versions are tested against. +# +# Pillow is pinned to 12.3.0 (same as requirements-synthid-scorer.txt) for +# the 24 known CVEs fixed after the upstream 9.4.0 pin. +torch==2.13.0 +transformers==5.15.0 +tokenizers==0.23.1 +datasets==5.0.1 +accelerate==1.14.0 +SentencePiece==0.2.2 +nltk==3.10.3 +jieba==0.42.1 +tqdm==4.70.0 +matplotlib==3.11.1 +Cython==3.2.9 +numpy==2.5.2 +scipy==1.18.0 +huggingface_hub==1.27.0 +Pillow==12.3.0 diff --git a/skills/remove-ai-marks/scripts/rewrite_text.py b/skills/remove-ai-marks/scripts/rewrite_text.py index e2946d0..afee474 100644 --- a/skills/remove-ai-marks/scripts/rewrite_text.py +++ b/skills/remove-ai-marks/scripts/rewrite_text.py @@ -26,9 +26,11 @@ import argparse import json import os import re +import subprocess import sys import urllib.error import urllib.request +from collections.abc import Callable from pathlib import Path from urllib.parse import urlparse @@ -37,6 +39,8 @@ sys.path.insert(0, str(Path(__file__).resolve().parent)) from common import cleaned_path, eprint, read_text_input, write_text_output # noqa: E402 from text_unicode import clean_text # noqa: E402 +DEFAULT_MARKLLM_MODEL = "facebook/opt-1.3b" + PROMPTS = { "paraphrase": ( "Rewrite the following text so that it uses substantially different wording at " @@ -169,6 +173,101 @@ class _NoRedirect(urllib.request.HTTPRedirectHandler): raise urllib.error.HTTPError(req.full_url, code, msg, headers, fp) +SCRIPTS_DIR = Path(__file__).resolve().parent + + +def _venv_python(upstream: Path) -> Path | None: + """Locate the MarkLLM checkout's venv interpreter, if it exists.""" + if os.name == "nt": + candidate = upstream / ".venv" / "Scripts" / "python.exe" + else: + candidate = upstream / ".venv" / "bin" / "python" + return candidate if candidate.is_file() else None + + +def _markllm_preexec() -> Callable[[], None] | None: + """Optional RLIMIT_AS guard for the MarkLLM child; None means "no limit". + + torch/CUDA usually needs large address spaces, so unlike the + exiftool/c2patool/SynthID children (common.subprocess_rlimits) this is + opt-in via WATERMARKS_MARKLLM_RLIMIT_AS (byte count, hex/octal allowed). + POSIX only; on Windows preexec_fn must stay None. + """ + raw = os.environ.get("WATERMARKS_MARKLLM_RLIMIT_AS") + if not raw or os.name != "posix": + return None + try: + limit = int(raw, 0) + except ValueError: + return None + + def _apply() -> None: + import resource + + resource.setrlimit(resource.RLIMIT_AS, (limit, limit)) + + return _apply + + +def _markllm_detect( + text: str, + *, + scheme: str, + upstream_dir: str, + model: str, + timeout: float, +) -> dict: + """Run the MarkLLM adapter on *text*; never fails the rewrite. + + Returns the adapter's JSON payload, or an ``available: False`` dict with + an ``error`` string when the backend is unconfigured or broken. The Layer B + rewrite proceeds regardless; MarkLLM verification is best-effort. + """ + if not upstream_dir: + return {"available": False, "error": "no MARKLLM_DIR set"} + upstream = Path(upstream_dir).expanduser().resolve() + if not upstream.is_dir() or not (upstream / "watermark").is_dir(): + return {"available": False, "error": f"MarkLLM checkout missing: {upstream}"} + venv_python = _venv_python(upstream) + if venv_python is None: + return {"available": False, "error": f"MarkLLM venv missing: {upstream}"} + + cmd = [ + str(venv_python), + str(SCRIPTS_DIR / "detect_text_watermark.py"), + "detect", + "-", + "--scheme", + scheme, + "--upstream-dir", + str(upstream), + "--model", + model, + "--json", + ] + try: + r = subprocess.run( + cmd, + input=text, + capture_output=True, + text=True, + timeout=timeout, + preexec_fn=_markllm_preexec(), + ) + except (OSError, subprocess.SubprocessError, TimeoutError) as e: + return {"available": False, "error": f"MarkLLM adapter error: {e}"} + + if r.returncode != 0: + return { + "available": False, + "error": (r.stderr or "").strip() or f"adapter exited {r.returncode}", + } + try: + return json.loads(r.stdout) + except ValueError as e: + return {"available": False, "error": f"adapter JSON parse error: {e}"} + + def build_prompt(strength: str, text: str, *, lang: str, original_lang: str) -> str: if strength == "paraphrase": return PROMPTS["paraphrase"].format(TEXT=text) @@ -276,6 +375,10 @@ def rewrite( temperature: float, candidates: int, allow_remote: bool = False, + markllm_scheme: str | None = None, + markllm_dir: str | None = None, + markllm_model: str | None = None, + markllm_timeout: float = 180.0, ) -> tuple[str, dict]: prompt = build_prompt(strength, text, lang=lang, original_lang=original_lang) info: dict = { @@ -288,6 +391,22 @@ def rewrite( "input_chars": len(text), } + markllm: dict | None = None + if markllm_scheme: + markllm = { + "scheme": markllm_scheme, + "before": _markllm_detect( + text, + scheme=markllm_scheme, + upstream_dir=markllm_dir or "", + model=markllm_model or DEFAULT_MARKLLM_MODEL, + timeout=markllm_timeout, + ), + } + if not markllm["before"]["available"]: + eprint(f"markllm verification unavailable: {markllm['before']['error']}") + info["markllm"] = markllm + if backend == "print-prompt": info["mode"] = "print-prompt" if candidates > 1: @@ -330,6 +449,26 @@ def rewrite( "Layer B is best-effort against statistical token-sampling watermarks; " "cannot certify removal against a vendor detector." ) + + if markllm: + after = _markllm_detect( + out, + scheme=markllm["scheme"], + upstream_dir=markllm_dir or "", + model=markllm_model or DEFAULT_MARKLLM_MODEL, + timeout=markllm_timeout, + ) + markllm["after"] = after + before = markllm["before"] + if before.get("available") and after.get("available"): + markllm["cleared"] = bool( + before.get("is_watermarked") and not after.get("is_watermarked") + ) + markllm["note"] = ( + "MarkLLM detection is only valid against the SAME scheme config + " + "keys used at generation; it does not certify a vendor detector." + ) + return out, info @@ -382,6 +521,29 @@ def main() -> int: help="Skip Layer A scrub on model output", ) p.add_argument("--json-stats", action="store_true", help="Stats JSON on stderr") + p.add_argument( + "--markllm-scheme", + choices=("kgw", "synthid", "synthid-text"), + default=None, + help="Optional: run MarkLLM before/after detection around the rewrite " + "(scheme = kgw or synthid)", + ) + p.add_argument( + "--markllm-dir", + default=_env("MARKLLM_DIR"), + help="MarkLLM checkout root (default: $MARKLLM_DIR)", + ) + p.add_argument( + "--markllm-model", + default=_env("MARKLLM_MODEL", DEFAULT_MARKLLM_MODEL), + help=f"Scoring model for MarkLLM detection (default: $MARKLLM_MODEL or {DEFAULT_MARKLLM_MODEL})", + ) + p.add_argument( + "--markllm-timeout", + type=float, + default=float(_env("WATERMARKS_MARKLLM_TIMEOUT", "180.0")), + help="Timeout per MarkLLM detection call (default: 180.0)", + ) p.add_argument( "--force-text", action="store_true", @@ -410,6 +572,10 @@ def main() -> int: temperature=args.temperature, candidates=args.candidates, allow_remote=allow_remote, + markllm_scheme=args.markllm_scheme, + markllm_dir=args.markllm_dir, + markllm_model=args.markllm_model, + markllm_timeout=args.markllm_timeout, ) except (urllib.error.URLError, TimeoutError, RuntimeError) as e: eprint(f"rewrite failed: {e}") diff --git a/skills/remove-ai-marks/scripts/setup_markllm.sh b/skills/remove-ai-marks/scripts/setup_markllm.sh new file mode 100755 index 0000000..89d3e9e --- /dev/null +++ b/skills/remove-ai-marks/scripts/setup_markllm.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Bootstrap an external THU-BPM/MarkLLM checkout for the optional text-watermark +# harness. +# +# The upstream project (https://github.com/THU-BPM/MarkLLM) is Apache-2.0 and +# is NOT bundled in this repository. This script clones it locally (pinned +# commit), creates a venv, and installs only the dependencies needed by +# detect_text_watermark.py (torch, transformers, datasets, ...). +# +# The base scoring model (default facebook/opt-1.3b) is downloaded from +# Hugging Face by detect_text_watermark.py at runtime, not here. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DEFAULT_DIR="${MARKLLM_DIR:-$HOME/MarkLLM}" +DIR="" +# Pinned upstream commit (2026-07-10). Do not point at a moving branch. +REF="c45ddc40f7b761beabe55a1b8dc4690e531d1c6d" +PYTHON="${PYTHON:-python3}" + +usage() { + cat <<'EOF' +Usage: setup_markllm.sh [--dir PATH] [--ref REF] [--python PYTHON] + +Clones (if needed) THU-BPM/MarkLLM, creates a venv, and installs the Python +dependencies required by detect_text_watermark.py (including torch). + +Options: + --dir PATH checkout directory (default: $MARKLLM_DIR or ~/MarkLLM) + --ref REF git ref to checkout (default: pinned commit SHA) + --python PY Python interpreter used to create the venv (default: python3) +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --dir) + DIR="${2:?--dir requires a value}" + shift 2 + ;; + --ref) + REF="${2:?--ref requires a value}" + shift 2 + ;; + --python) + PYTHON="${2:?--python requires a value}" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "unknown option: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +DIR="${DIR:-$DEFAULT_DIR}" +mkdir -p "$(dirname "$DIR")" +if command -v realpath >/dev/null 2>&1; then + DIR="$(realpath -m "$DIR")" +else + DIR="$(cd "$(dirname "$DIR")" && pwd)/$(basename "$DIR")" +fi + +if [[ ! -d "$DIR/.git" ]]; then + echo "Cloning THU-BPM/MarkLLM into $DIR (pinned ref: $REF)" + git clone --depth 1 --filter=blob:none --sparse \ + https://github.com/THU-BPM/MarkLLM.git "$DIR" + git -C "$DIR" fetch --depth 1 origin "$REF" + git -C "$DIR" checkout --detach "$REF" + git -C "$DIR" sparse-checkout set --no-cone \ + '/watermark/' \ + '/config/' \ + '/utils/' \ + '/exceptions/' \ + '/evaluation/dataset.py' \ + '/LICENSE' \ + '/README.md' + HEAD_SHA="$(git -C "$DIR" rev-parse HEAD)" + if [[ "$HEAD_SHA" != "$REF" ]]; then + echo "error: expected pinned ref $REF, got $HEAD_SHA" >&2 + exit 1 + fi +else + echo "Using existing checkout: $DIR" +fi + +if [[ ! -x "$DIR/.venv/bin/python" ]]; then + echo "Creating venv at $DIR/.venv" + "$PYTHON" -m venv "$DIR/.venv" +fi + +echo "Installing Python dependencies" +# Pin pip itself (unpinned --upgrade pip was a supply-chain drift point). +"$DIR/.venv/bin/python" -m pip install --upgrade "pip==26.2.1" + +# Install torch with the right platform index before the other pinned deps. +if command -v nvidia-smi >/dev/null 2>&1; then + cuda="$(nvidia-smi 2>/dev/null | sed -n 's/.*CUDA Version: \([0-9]*\.[0-9]*\).*/\1/p' | head -1)" + if [[ -n "$cuda" ]]; then + tag="cu${cuda/./}" + index="https://download.pytorch.org/whl/$tag" + echo "NVIDIA GPU detected (CUDA $cuda); installing torch from $index" + "$DIR/.venv/bin/python" -m pip install torch --index-url "$index" + else + echo "nvidia-smi present but no CUDA version found; installing default torch" + "$DIR/.venv/bin/python" -m pip install torch + fi +else + echo "No NVIDIA GPU detected; installing default torch (CPU/MPS)" + "$DIR/.venv/bin/python" -m pip install torch +fi + +"$DIR/.venv/bin/python" -m pip install -r "$SCRIPT_DIR/requirements-markllm.txt" + +cat < str: + detect_body = 'raise RuntimeError("boom")' if fail_detect else 'return {"is_watermarked": True, "score": 3.5}' + gen_body = 'raise RuntimeError("boom")' if fail_generate else "return 'WATERMARKED SAMPLE'" + return ( + "class _WM:\n" + " def __init__(self):\n" + " self.config = SimpleNamespace(gen_kwargs={})\n" + " def detect_watermark(self, text, return_dict=True):\n" + f" {detect_body}\n" + " def generate_watermarked_text(self, prompt):\n" + f" {gen_body}\n" + " def generate_unwatermarked_text(self, prompt):\n" + " return 'PLAIN SAMPLE'\n" + "\n" + "class AutoWatermark:\n" + " @staticmethod\n" + " def load(algorithm_name, algorithm_config=None, transformers_config=None):\n" + " return _WM()\n" + ) + + +def _make_fake_upstream( + tmp_path: Path, + *, + with_config: bool = True, + fail_detect: bool = False, + fail_generate: bool = False, + missing_watermark_dir: bool = False, +) -> Path: + upstream = tmp_path / "MarkLLM" + config_dir = upstream / "config" + config_dir.mkdir(parents=True) + if with_config: + (config_dir / "KGW.json").write_text(KGW_CONFIG) + (config_dir / "SynthID.json").write_text(SYNTHID_CONFIG) + if not missing_watermark_dir: + watermark = upstream / "watermark" + watermark.mkdir(parents=True) + (watermark / "__init__.py").write_text("") + (watermark / "auto_watermark.py").write_text( + "from types import SimpleNamespace\n" + _fake_auto_watermark( + fail_detect=fail_detect, fail_generate=fail_generate + ) + ) + utils_dir = upstream / "utils" + utils_dir.mkdir(parents=True) + (utils_dir / "__init__.py").write_text("") + (utils_dir / "transformers_config.py").write_text(FAKE_TRANSFORMERS_CONFIG) + transformers_dir = upstream / "transformers" + transformers_dir.mkdir(parents=True) + (transformers_dir / "__init__.py").write_text(FAKE_TRANSFORMERS) + return upstream + + +def _run_adapter(*args: str) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env.pop("MARKLLM_DIR", None) + return subprocess.run( + [sys.executable, str(DETECT_SCRIPT), *args], + capture_output=True, + text=True, + env=env, + ) + + +def test_cli_unavailable_without_upstream(tmp_path: Path): + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter("detect", str(f), "--scheme", "kgw") + assert r.returncode == 3 + assert "MARKLLM_DIR" in (r.stderr or "") + + +def test_cli_unavailable_incomplete_checkout(tmp_path: Path): + f = tmp_path / "t.txt" + f.write_text("hello world") + empty = tmp_path / "empty" + empty.mkdir() + r = _run_adapter("detect", str(f), "--scheme", "kgw", "--upstream-dir", str(empty)) + assert r.returncode == 3 + + upstream = _make_fake_upstream(tmp_path, missing_watermark_dir=True) + r = _run_adapter("detect", str(f), "--scheme", "kgw", "--upstream-dir", str(upstream)) + assert r.returncode == 3 + + +def test_cli_unavailable_missing_config(tmp_path: Path): + f = tmp_path / "t.txt" + f.write_text("hello world") + upstream = _make_fake_upstream(tmp_path, with_config=False) + r = _run_adapter("detect", str(f), "--scheme", "kgw", "--upstream-dir", str(upstream)) + assert r.returncode == 3 + assert "config" in (r.stderr or "").lower() + + +def test_cli_unavailable_missing_deps(tmp_path: Path): + # The watermark module imports a nonexistent dependency -> ImportError -> + # exit 3 ("dependencies missing") before any model download. + upstream = tmp_path / "MarkLLM" + (upstream / "config").mkdir(parents=True) + (upstream / "config" / "KGW.json").write_text(KGW_CONFIG) + watermark = upstream / "watermark" + watermark.mkdir() + (watermark / "__init__.py").write_text("") + (watermark / "auto_watermark.py").write_text("import does_not_exist_123\n") + (upstream / "utils").mkdir() + (upstream / "utils" / "__init__.py").write_text("") + (upstream / "utils" / "transformers_config.py").write_text(FAKE_TRANSFORMERS_CONFIG) + (upstream / "transformers").mkdir() + (upstream / "transformers" / "__init__.py").write_text(FAKE_TRANSFORMERS) + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter("detect", str(f), "--scheme", "kgw", "--upstream-dir", str(upstream)) + assert r.returncode == 3 + assert "dependencies missing" in (r.stderr or "") + + +def test_cli_bad_input_missing_file(tmp_path: Path): + r = _run_adapter("detect", str(tmp_path / "missing.txt"), "--scheme", "kgw") + assert r.returncode == 2 + + +def test_cli_bad_input_binary(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path) + png = tmp_path / "img.png" + png.write_bytes(b"\x89PNG\r\n\x1a\nnot really") + r = _run_adapter("detect", str(png), "--scheme", "kgw", "--upstream-dir", str(upstream)) + assert r.returncode == 2 + assert "refusing" in (r.stderr or "") + + +def test_cli_bad_scheme(tmp_path: Path): + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter("detect", str(f), "--scheme", "nope") + assert r.returncode == 2 + + +def test_cli_detect_json_success(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path) + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter( + "detect", str(f), "--scheme", "kgw", "--upstream-dir", str(upstream), + "--device", "cpu", "--json", + ) + assert r.returncode == 0, r.stderr + payload = json.loads(r.stdout) + assert payload["available"] is True + assert payload["scheme"] == "KGW" + assert payload["is_watermarked"] is True + assert payload["score"] == 3.5 + assert payload["threshold"] == 4.0 + assert payload["device"] == "cpu" + + +def test_cli_detect_synthid_alias(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path) + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter( + "detect", str(f), "--scheme", "synthid-text", "--upstream-dir", str(upstream), + "--device", "cpu", "--json", + ) + assert r.returncode == 0, r.stderr + payload = json.loads(r.stdout) + assert payload["scheme"] == "SynthID" + assert payload["threshold"] == 0.52 + + +def test_cli_detect_runtime_error(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path, fail_detect=True) + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter( + "detect", str(f), "--scheme", "kgw", "--upstream-dir", str(upstream), + "--device", "cpu", "--json", + ) + assert r.returncode == 1 + assert "boom" in (r.stderr or "") + + +def test_cli_detect_offline_flag(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path) + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter( + "detect", str(f), "--scheme", "kgw", "--upstream-dir", str(upstream), + "--device", "cpu", "--json", "--offline", + ) + assert r.returncode == 0, r.stderr + assert "local_files_only" in (r.stderr or "") + assert "True" in (r.stderr or "") + + +def test_cli_config_too_large(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path) + big = tmp_path / "huge.json" + big.write_bytes(b"x" * (1024 * 1024 + 1)) + f = tmp_path / "t.txt" + f.write_text("hello world") + r = _run_adapter( + "detect", str(f), "--scheme", "kgw", "--config", str(big), + "--upstream-dir", str(upstream), + ) + assert r.returncode == 3 + assert "too large" in (r.stderr or "") + + +def test_cli_watermark_json_success(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path) + prompt = tmp_path / "prompt.txt" + prompt.write_text("write about capybaras") + wm_out = tmp_path / "wm.txt" + uwm_out = tmp_path / "uwm.txt" + r = _run_adapter( + "watermark", str(prompt), "--scheme", "kgw", + "-o", str(wm_out), "-o2", str(uwm_out), + "--upstream-dir", str(upstream), "--device", "cpu", "--json", + ) + assert r.returncode == 0, r.stderr + payload = json.loads(r.stdout) + assert payload["available"] is True + assert wm_out.read_text() == "WATERMARKED SAMPLE" + assert uwm_out.read_text() == "PLAIN SAMPLE" + + +def test_cli_watermark_runtime_error(tmp_path: Path): + upstream = _make_fake_upstream(tmp_path, fail_generate=True) + prompt = tmp_path / "prompt.txt" + prompt.write_text("write about capybaras") + r = _run_adapter( + "watermark", str(prompt), "--scheme", "kgw", + "--upstream-dir", str(upstream), "--device", "cpu", "--json", + ) + assert r.returncode == 1 + assert "boom" in (r.stderr or "") + + +def test_rewrite_markllm_detect_missing_venv(tmp_path: Path): + import rewrite_text + + upstream = tmp_path / "MarkLLM" + upstream.mkdir() + result = rewrite_text._markllm_detect( + "hello", scheme="kgw", upstream_dir=str(upstream), model="x", timeout=5, + ) + assert result["available"] is False + + +def test_rewrite_markllm_detect_parses_json(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + import rewrite_text + + upstream = tmp_path / "MarkLLM" + if os.name == "nt": + venv_python = upstream / ".venv" / "Scripts" / "python.exe" + else: + venv_python = upstream / ".venv" / "bin" / "python" + venv_python.parent.mkdir(parents=True) + venv_python.write_text("") + (upstream / "watermark").mkdir() + + payload = {"available": True, "is_watermarked": True, "score": 2.0} + captured: dict = {} + + def fake_run(cmd, **kwargs): + captured["cmd"] = cmd + captured["input"] = kwargs.get("input") + return SimpleNamespace(returncode=0, stdout=json.dumps(payload), stderr="") + + monkeypatch.setattr(rewrite_text.subprocess, "run", fake_run) + result = rewrite_text._markllm_detect( + "hello", scheme="kgw", upstream_dir=str(upstream), model="x", timeout=5, + ) + assert result["available"] is True + assert result["score"] == 2.0 + assert captured["input"] == "hello" + assert captured["cmd"][0] == str(venv_python) + + +def test_rewrite_markllm_detect_adapter_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + import rewrite_text + + upstream = tmp_path / "MarkLLM" + if os.name == "nt": + venv_python = upstream / ".venv" / "Scripts" / "python.exe" + else: + venv_python = upstream / ".venv" / "bin" / "python" + venv_python.parent.mkdir(parents=True) + venv_python.write_text("") + (upstream / "watermark").mkdir() + + def fake_run(cmd, **kwargs): + return SimpleNamespace(returncode=3, stdout="", stderr="deps missing") + + monkeypatch.setattr(rewrite_text.subprocess, "run", fake_run) + result = rewrite_text._markllm_detect( + "hello", scheme="kgw", upstream_dir=str(upstream), model="x", timeout=5, + ) + assert result["available"] is False + assert "deps missing" in result["error"] + + +def test_markllm_preexec_default_off(monkeypatch: pytest.MonkeyPatch): + import rewrite_text + + monkeypatch.delenv("WATERMARKS_MARKLLM_RLIMIT_AS", raising=False) + assert rewrite_text._markllm_preexec() is None + + +def test_markllm_preexec_env(monkeypatch: pytest.MonkeyPatch): + import rewrite_text + + if os.name != "posix": + pytest.skip("preexec_fn is POSIX-only") + monkeypatch.setenv("WATERMARKS_MARKLLM_RLIMIT_AS", "0x40000000") + fn = rewrite_text._markllm_preexec() + assert callable(fn) + + +def test_rewrite_markllm_detect_applies_rlimit( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +): + import rewrite_text + + if os.name != "posix": + pytest.skip("preexec_fn is POSIX-only") + upstream = tmp_path / "MarkLLM" + venv_python = upstream / ".venv" / "bin" / "python" + venv_python.parent.mkdir(parents=True) + venv_python.write_text("") + (upstream / "watermark").mkdir() + monkeypatch.setenv("WATERMARKS_MARKLLM_RLIMIT_AS", "1073741824") + captured: dict = {} + + def fake_run(cmd, **kwargs): + captured["preexec_fn"] = kwargs.get("preexec_fn") + return SimpleNamespace( + returncode=0, stdout='{"available": true}', stderr="" + ) + + monkeypatch.setattr(rewrite_text.subprocess, "run", fake_run) + result = rewrite_text._markllm_detect( + "hello", scheme="kgw", upstream_dir=str(upstream), model="x", timeout=5, + ) + assert result["available"] is True + assert callable(captured["preexec_fn"]) + + +def test_rewrite_markllm_hook_records_before_after(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + import rewrite_text + + def fake_detect(text, **kwargs): + return {"available": True, "is_watermarked": text == "ORIG", "score": 3.0} + + monkeypatch.setattr(rewrite_text, "_markllm_detect", fake_detect) + monkeypatch.setattr( + rewrite_text, "call_ollama", lambda *a, **k: "REWRITTEN OUTPUT" + ) + out, info = rewrite_text.rewrite( + "ORIG", + backend="ollama", + model="m", + base_url="http://127.0.0.1:11434", + api_key=None, + strength="paraphrase", + lang="French", + original_lang="English", + timeout=10, + layer_a_after=False, + temperature=0.9, + candidates=1, + markllm_scheme="kgw", + markllm_dir="/tmp/x", + markllm_model="opt-1.3b", + markllm_timeout=5, + ) + assert out == "REWRITTEN OUTPUT" + mk = info["markllm"] + assert mk["before"]["is_watermarked"] is True + assert mk["after"]["is_watermarked"] is False + assert mk["cleared"] is True + assert "note" in mk