Detectie van watermerken in LLM-output

"""
LLM-watermerkdetectie-implementatie voor forensische analyse.
 
Gebaseerd op het green-list-schema van Kirchenbauer et al. (2023).
Deze module implementeert de detectiekant -- gegeven tekst en de
watermerksleutel, bepalen of de tekst was voorzien van een watermerk.
"""
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]:
    """
    Bereken de green list voor een gegeven contexttoken.
 
    Gebruikt HMAC-SHA256 om de woordenschat deterministisch te verdelen
    in green- en red-lists op basis van het vorige token en de geheime sleutel.
 
    Args:
        previous_token_id: De token-ID die aan de huidige positie voorafgaat.
        vocab_size: Totale grootte van de woordenschat.
        secret_key: De geheime watermerksleutel.
        gamma: Fractie van de woordenschat in de green list (standaard 0.5).
 
    Returns:
        Set van token-ID's in de green list.
    """
    # Maak deterministische seed van vorige token en geheime sleutel
    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 of een reeks token-ID's een watermerk bevat.
 
    Voert een z-toets voor één proportie uit tegen de nulhypothese
    dat tokens zonder green-list-bias worden getrokken.
 
    Onder de nulhypothese (geen watermerk) is de verwachte fractie van
    green-list-tokens gamma. Onder het alternatief (van watermerk voorzien)
    zou de fractie significant hoger moeten zijn.
 
    Args:
        token_ids: Reeks te testen token-ID's.
        vocab_size: Totale grootte van de woordenschat.
        secret_key: De geheime watermerksleutel.
        gamma: Green-list-fractie gebruikt tijdens het watermerken.
 
    Returns:
        Dict met detectieresultaten inclusief z-score en p-waarde.
    """
    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
 
    # Z-toets voor één proportie
    # H0: p = gamma (geen watermerk)
    # H1: p > gamma (watermerk aanwezig)
    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]]:
    """
    Tokeniseer tekst met meerdere kandidaat-modeltokenizers.
 
    Omdat het watermerk op tokenniveau is ingebed, moeten we
    testen met de tokenizer van elk kandidaatmodel. De juiste
    tokenizer produceert een sterker watermerksignaal.
    """
    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 tekst op watermerken van meerdere kandidaatmodellen/-providers.
 
    Test elke combinatie van tokenizer en watermerksleutel om
    te bepalen welke (indien aanwezig) de tekst produceerde.
    """
    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]:
    """
    Pas watermerkdetectie toe in een sliding window over de tekst.
 
    Dit identificeert welke delen van een document waarschijnlijk door AI
    zijn gegenereerd versus door mensen geschreven, op basis van de lokale
    sterkte van het watermerksignaal.
    """
    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:
    """
    Analyseer tekst op tekenen van pogingen tot watermerkverwijdering.
 
    Zelfs na parafraseren of bewerken kunnen statistische artefacten van
    het oorspronkelijke watermerk achterblijven in delen van de tekst
    die niet zijn gewijzigd.
    """
    # Detectie over volledige tekst
    full_result = detect_watermark(token_ids, vocab_size, secret_key, gamma)
 
    # Geveensterde detectie om overlevende watermerkfragmenten te vinden
    window_results = windowed_watermark_detection(
        token_ids, vocab_size, secret_key,
        window_size=50, step_size=10, gamma=gamma,
    )
 
    # Tel vensters met significant watermerksignaal
    significant_windows = [
        w for w in window_results if w.get("z_score", 0) > 2.0
    ]
 
    # Analyseer het patroon van overlevende vs. verwijderde watermerkregio's
    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"