LLM Output Watermark Detection

"""
LLM watermark 偵測 實作 for forensic analysis.
 
Based on the green-list scheme from Kirchenbauer et al. (2023).
This module implements the 偵測 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 符元.
 
    Uses HMAC-SHA256 to deterministically partition the vocabulary
    into green and red lists based on the previous 符元 and secret key.
 
    Args:
        previous_token_id: The 符元 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 符元 IDs in the green list.
    """
    # Create deterministic seed from previous 符元 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:
    """
    測試 whether a sequence of 符元 IDs contains a watermark.
 
    Performs a one-proportion z-測試 against the null hypothesis
    that 符元 are drawn without green-list bias.
 
    Under the null (no watermark), the expected fraction of green-list
    符元 is gamma. Under the alternative (watermarked), the fraction
    should be significantly higher.
 
    Args:
        token_ids: Sequence of 符元 IDs to 測試.
        vocab_size: Total vocabulary size.
        secret_key: The watermarking secret key.
        gamma: Green list fraction used during watermarking.
 
    Returns:
        Dict with 偵測 results including z-score and p-value.
    """
    if len(token_ids) < 2:
        return {"error": "Need at least 2 符元 for 偵測"}
 
    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-測試
    # 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 符元 level, we must
    測試 with each candidate model's 分詞器. The correct
    分詞器 will produce a stronger watermark signal.
    """
    results = {}
    for model_name in model_candidates:
        分詞器 = AutoTokenizer.from_pretrained(model_name)
        token_ids = 分詞器.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 分詞器 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
 
        分詞器 = AutoTokenizer.from_pretrained(model_name)
        vocab_size = 分詞器.vocab_size
 
        偵測 = detect_watermark(
            token_ids=token_ids,
            vocab_size=vocab_size,
            secret_key=candidate_keys[model_name],
            gamma=gamma,
        )
        偵測["model"] = model_name
        results.append(偵測)
 
    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 偵測 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]
        偵測 = detect_watermark(window, vocab_size, secret_key, gamma)
        偵測["token_start"] = start
        偵測["token_end"] = start + window_size
        results.append(偵測)
 
    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 偵測
    full_result = detect_watermark(token_ids, vocab_size, secret_key, gamma)
 
    # Windowed 偵測 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"