Watermarking & AI-Generated Text Detection

import hashlib
import numpy as np
 
class WatermarkLogitsProcessor:
    """Simplified watermark injection during LLM inference."""
 
    def __init__(self, secret_key: bytes, gamma: float = 0.5, delta: float = 2.0):
        self.secret_key = secret_key
        self.gamma = gamma  # fraction of vocab in green list
        self.delta = delta  # logit bias added to green-list tokens
 
    def get_green_list(self, prev_token_id: int, vocab_size: int) -> set:
        """Deterministically partition vocab into green/red lists."""
        seed = hashlib.sha256(
            self.secret_key + prev_token_id.to_bytes(4, "big")
        ).digest()
        rng = np.random.RandomState(
            int.from_bytes(seed[:4], "big")
        )
        green_count = int(vocab_size * self.gamma)
        green_list = set(rng.choice(vocab_size, green_count, replace=False))
        return green_list
 
    def __call__(self, prev_token_id: int, logits: np.ndarray) -> np.ndarray:
        """Add delta bias to green-list token logits."""
        green_list = self.get_green_list(prev_token_id, len(logits))
        modified = logits.copy()
        for token_id in green_list:
            modified[token_id] += self.delta
        return modified

from scipy import stats
 
def detect_watermark(
    text_token_ids: list[int],
    secret_key: bytes,
    vocab_size: int,
    gamma: float = 0.5,
) -> dict:
    """Test whether text contains a statistical watermark."""
    green_count = 0
    total = len(text_token_ids) - 1  # skip first token (no predecessor)
 
    for i in range(1, len(text_token_ids)):
        prev_id = text_token_ids[i - 1]
        green_list = get_green_list(prev_id, vocab_size, secret_key, gamma)
        if text_token_ids[i] in green_list:
            green_count += 1
 
    # Under null hypothesis (no watermark), green fraction ~ gamma
    z_score = (green_count - gamma * total) / np.sqrt(total * gamma * (1 - gamma))
    p_value = 1 - stats.norm.cdf(z_score)
 
    return {
        "green_fraction": green_count / total,
        "z_score": z_score,
        "p_value": p_value,
        "watermarked": z_score > 4.0,  # threshold
    }

def detect_gpt_score(
    text: str,
    scoring_model,
    mask_model,
    n_perturbations: int = 100,
) -> float:
    """
    DetectGPT: AI text sits at local maxima of log-probability.
    Perturbations should decrease log-prob for AI text,
    but have random effect on human text.
    """
    original_ll = scoring_model.log_likelihood(text)
    perturbation_lls = []
 
    for _ in range(n_perturbations):
        # Mask random spans and fill with mask_model (e.g., T5)
        perturbed = mask_and_fill(text, mask_model)
        perturbation_lls.append(scoring_model.log_likelihood(perturbed))
 
    mean_perturbed_ll = np.mean(perturbation_lls)
    std_perturbed_ll = np.std(perturbation_lls)
 
    # Normalized perturbation discrepancy
    score = (original_ll - mean_perturbed_ll) / std_perturbed_ll
    return score  # high score -> likely AI-generated

def paraphrase_attack(watermarked_text: str, paraphraser) -> str:
    """Use a separate model to remove watermark via paraphrasing."""
    prompt = f"""Rewrite the following text to convey the same meaning
using different words and sentence structures. Preserve all
factual content but change the phrasing completely.
 
Text: {watermarked_text}
 
Rewritten:"""
    return paraphraser.generate(prompt)

def synonym_substitution_attack(
    text: str,
    substitution_rate: float = 0.25,
    synonym_dict: dict = None,
) -> str:
    """Replace ~25% of content words with synonyms."""
    words = text.split()
    for i, word in enumerate(words):
        if random.random() < substitution_rate and word.lower() in synonym_dict:
            words[i] = random.choice(synonym_dict[word.lower()])
    return " ".join(words)

ZERO_WIDTH_CHARS = ["\u200b", "\u200c", "\u200d", "\ufeff"]
 
def unicode_injection_attack(text: str, injection_rate: float = 0.15) -> str:
    """Insert zero-width characters to disrupt token-chain watermark."""
    chars = list(text)
    positions = random.sample(range(len(chars)), int(len(chars) * injection_rate))
    for pos in sorted(positions, reverse=True):
        chars.insert(pos, random.choice(ZERO_WIDTH_CHARS))
    return "".join(chars)

def cross_model_transfer(watermarked_text: str, clean_model) -> str:
    """Transfer content through a non-watermarked model."""
    prompt = f"""Based on the following information, write a comprehensive
response covering the same topics and conclusions:
 
{watermarked_text}"""
    return clean_model.generate(prompt)