LLM Output Watermark Detection

"""
LLM watermark detection implementation for forensic analysis.
 
Based on the green-list scheme from Kirchenbauer et al. (2023).
This module implements the detection side -- given text and the
watermarking key, determine whether the text was watermarked.
"""
import hashlib
from collections.abc import Sequence
 
import numpy as np
from scipy import stats
 
 
def compute_green_list(
    previous_token_id: int,
    vocab_size: int,
    secret_key: bytes,
    gamma: float = 0.5,
) -> set[int]:
    """
    Compute the green list for a given context token.
 
    Uses HMAC-SHA256 to deterministically partition the vocabulary
    into green and red lists based on the previous token and secret key.
 
    Args:
        previous_token_id: The token ID preceding the current position.
        vocab_size: Total vocabulary size.
        secret_key: The watermarking secret key.
        gamma: Fraction of vocabulary in the green list (default 0.5).
 
    Returns:
        Set of token IDs in the green list.
    """
    # Create deterministic seed from previous token and secret key
    seed_material = previous_token_id.to_bytes(4, 'big') + secret_key
    hash_bytes = hashlib.sha256(seed_material).digest()
    seed = int.from_bytes(hash_bytes[:8], 'big')
 
    rng = np.random.Generator(np.random.PCG64(seed))
    permutation = rng.permutation(vocab_size)
 
    green_list_size = int(gamma * vocab_size)
    return set(permutation[:green_list_size].tolist())
 
 
def detect_watermark(
    token_ids: Sequence[int],
    vocab_size: int,
    secret_key: bytes,
    gamma: float = 0.5,
) -> dict:
    """
    Test whether a sequence of token IDs contains a watermark.
 
    Performs a one-proportion z-test against the null hypothesis
    that tokens are drawn without green-list bias.
 
    Under the null (no watermark), the expected fraction of green-list
    tokens is gamma. Under the alternative (watermarked), the fraction
    should be significantly higher.
 
    Args:
        token_ids: Sequence of token IDs to test.
        vocab_size: Total vocabulary size.
        secret_key: The watermarking secret key.
        gamma: Green list fraction used during watermarking.
 
    Returns:
        Dict with detection results including z-score and p-value.
    """
    if len(token_ids) < 2:
        return {"error": "Need at least 2 tokens for detection"}
 
    green_count = 0
    total_tested = 0
 
    for i in range(1, len(token_ids)):
        prev_token = token_ids[i - 1]
        current_token = token_ids[i]
        green_list = compute_green_list(prev_token, vocab_size, secret_key, gamma)
 
        if current_token in green_list:
            green_count += 1
        total_tested += 1
 
    observed_fraction = green_count / total_tested
    expected_fraction = gamma
 
    # One-proportion z-test
    # H0: p = gamma (no watermark)
    # H1: p > gamma (watermark present)
    se = np.sqrt(gamma * (1 - gamma) / total_tested)
    z_score = (observed_fraction - expected_fraction) / se
    p_value = 1 - stats.norm.cdf(z_score)
 
    return {
        "total_tokens_tested": total_tested,
        "green_list_count": green_count,
        "observed_green_fraction": round(observed_fraction, 4),
        "expected_green_fraction": gamma,
        "z_score": round(float(z_score), 4),
        "p_value": float(p_value),
        "watermark_detected": p_value < 1e-5,
        "confidence": "high" if p_value < 1e-10 else "medium" if p_value < 1e-5 else "low",
    }

from transformers import AutoTokenizer
 
 
def prepare_text_for_detection(
    text: str,
    model_candidates: list[str],
) -> dict[str, list[int]]:
    """
    Tokenize text with multiple candidate model tokenizers.
 
    Since the watermark is embedded at the token level, we must
    test with each candidate model's tokenizer. The correct
    tokenizer will produce a stronger watermark signal.
    """
    results = {}
    for model_name in model_candidates:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        token_ids = tokenizer.encode(text, add_special_tokens=False)
        results[model_name] = token_ids
    return results
 
 
def multi_model_watermark_scan(
    text: str,
    model_candidates: list[str],
    candidate_keys: dict[str, bytes],
    gamma: float = 0.5,
) -> list[dict]:
    """
    Scan text for watermarks from multiple candidate models/providers.
 
    Tests each combination of tokenizer and watermarking key to
    determine which (if any) produced the text.
    """
    tokenizations = prepare_text_for_detection(text, model_candidates)
    results = []
 
    for model_name, token_ids in tokenizations.items():
        if model_name not in candidate_keys:
            continue
 
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        vocab_size = tokenizer.vocab_size
 
        detection = detect_watermark(
            token_ids=token_ids,
            vocab_size=vocab_size,
            secret_key=candidate_keys[model_name],
            gamma=gamma,
        )
        detection["model"] = model_name
        results.append(detection)
 
    results.sort(key=lambda r: r.get("z_score", 0), reverse=True)
    return results

def windowed_watermark_detection(
    token_ids: list[int],
    vocab_size: int,
    secret_key: bytes,
    window_size: int = 100,
    step_size: int = 25,
    gamma: float = 0.5,
) -> list[dict]:
    """
    Apply watermark detection in a sliding window across the text.
 
    This identifies which portions of a document are likely AI-generated
    versus human-written, based on local watermark signal strength.
    """
    results = []
 
    for start in range(0, len(token_ids) - window_size, step_size):
        window = token_ids[start:start + window_size]
        detection = detect_watermark(window, vocab_size, secret_key, gamma)
        detection["token_start"] = start
        detection["token_end"] = start + window_size
        results.append(detection)
 
    return results

def analyze_removal_artifacts(
    token_ids: list[int],
    vocab_size: int,
    secret_key: bytes,
    gamma: float = 0.5,
) -> dict:
    """
    Analyze text for signs of watermark removal attempts.
 
    Even after paraphrasing or editing, statistical artifacts of
    the original watermark may remain in portions of the text
    that were not modified.
    """
    # Full-text detection
    full_result = detect_watermark(token_ids, vocab_size, secret_key, gamma)
 
    # Windowed detection to find surviving watermark fragments
    window_results = windowed_watermark_detection(
        token_ids, vocab_size, secret_key,
        window_size=50, step_size=10, gamma=gamma,
    )
 
    # Count windows with significant watermark signal
    significant_windows = [
        w for w in window_results if w.get("z_score", 0) > 2.0
    ]
 
    # Analyze the pattern of surviving vs. removed watermark regions
    z_scores = [w.get("z_score", 0) for w in window_results]
    z_variance = float(np.var(z_scores)) if z_scores else 0.0
 
    return {
        "full_text_detection": full_result,
        "total_windows": len(window_results),
        "significant_windows": len(significant_windows),
        "significant_fraction": (
            len(significant_windows) / len(window_results)
            if window_results else 0.0
        ),
        "z_score_variance": z_variance,
        "removal_likely": (
            full_result.get("z_score", 0) < 2.0
            and len(significant_windows) > len(window_results) * 0.2
        ),
        "interpretation": _interpret_removal_analysis(
            full_result, significant_windows, window_results, z_variance
        ),
    }
 
 
def _interpret_removal_analysis(
    full_result: dict,
    significant_windows: list,
    all_windows: list,
    z_variance: float,
) -> str:
    full_z = full_result.get("z_score", 0)
    sig_frac = len(significant_windows) / max(len(all_windows), 1)
 
    if full_z > 4.0:
        return "Strong watermark detected -- no significant removal"
    if full_z > 2.0 and sig_frac > 0.5:
        return "Partial watermark -- possible light editing or paraphrasing of some sections"
    if full_z < 2.0 and sig_frac > 0.2:
        return "Watermark largely removed but fragments survive -- likely systematic removal attempt"
    if z_variance > 4.0:
        return "High variance in local watermark strength -- mixed human/AI or selective editing"
    return "No watermark signal detected"