浮水印移除技術

import hashlib
import torch
 
def watermarked_sampling(logits, prev_token, watermark_key, bias=2.0):
    """取樣時對 'green list' token 加偏誤。
    green list 由前一 token 與浮水印金鑰之 hash 決定。"""
    # 依 前一 token + 金鑰 之確定性 green/red 分割
    seed = int(hashlib.sha256(
        f"{watermark_key}:{prev_token}".encode()
    ).hexdigest()[:8], 16)
    rng = torch.Generator().manual_seed(seed)
 
    vocab_size = logits.shape[-1]
    perm = torch.randperm(vocab_size, generator=rng)
    green_tokens = perm[:vocab_size // 2]
 
    # 對 logits 向 green tokens 加偏誤
    watermarked_logits = logits.clone()
    watermarked_logits[green_tokens] += bias
 
    return watermarked_logits
 
def detect_watermark(text, tokenizer, watermark_key, threshold=0.5):
    """藉由測量 green token 比例偵測浮水印。"""
    tokens = tokenizer.encode(text)
    green_count = 0
 
    for i in range(1, len(tokens)):
        seed = int(hashlib.sha256(
            f"{watermark_key}:{tokens[i-1]}".encode()
        ).hexdigest()[:8], 16)
        rng = torch.Generator().manual_seed(seed)
        perm = torch.randperm(tokenizer.vocab_size, generator=rng)
        green_tokens = set(perm[:tokenizer.vocab_size // 2].tolist())
 
        if tokens[i] in green_tokens:
            green_count += 1
 
    green_fraction = green_count / (len(tokens) - 1)
    # 無浮水印文字：~50% green tokens
    # 有浮水印文字：顯著 > 50% green tokens
    z_score = (green_fraction - 0.5) / (0.5 / (len(tokens) - 1) ** 0.5)
    return z_score > threshold, z_score, green_fraction

def paraphrase_removal(watermarked_text, paraphrase_model, tokenizer):
    """以未加浮水印之模型換句話以移除浮水印。"""
    prompt = f"""Rewrite the following text to convey the same meaning
using different words and sentence structures. Preserve all factual
content and technical accuracy.
 
Original text:
{watermarked_text}
 
Rewritten text:"""
 
    # 換句話模型產生不含浮水印偏誤之新 token
    paraphrased = paraphrase_model.generate(
        tokenizer(prompt, return_tensors="pt").input_ids,
        max_new_tokens=len(watermarked_text.split()) * 2,
        temperature=0.7,
    )
    return tokenizer.decode(paraphrased[0], skip_special_tokens=True)

def token_substitution_removal(text, tokenizer, watermark_key, model):
    """以語意相當之 red token 選擇性替換 green token
    以移除浮水印。"""
    tokens = tokenizer.encode(text)
    modified_tokens = list(tokens)
 
    for i in range(1, len(tokens)):
        # 判斷當前 token 是否為 green
        green_set = get_green_tokens(tokens[i-1], watermark_key, tokenizer.vocab_size)
 
        if tokens[i] in green_set:
            # 尋找語意相近之 red token
            # 取得該位置上模型之機率分布
            context = torch.tensor([modified_tokens[:i]])
            logits = model(context).logits[0, -1]
 
            # 尋找最高機率之 red token
            red_set = set(range(tokenizer.vocab_size)) - green_set
            red_probs = [(t, logits[t].item()) for t in red_set]
            red_probs.sort(key=lambda x: x[1], reverse=True)
 
            best_red_token = red_probs[0][0]
            # 僅在 red token 機率合理時才替換
            if logits[best_red_token] > logits[tokens[i]] - 3.0:
                modified_tokens[i] = best_red_token
 
    return tokenizer.decode(modified_tokens)

def embedding_perturbation_removal(text, model, tokenizer, noise_scale=0.1):
    """藉由對 token embedding 加雜訊並重新解碼，
    移除 embedding 層級之浮水印。"""
    tokens = tokenizer(text, return_tensors="pt")
    embeddings = model.get_input_embeddings()(tokens.input_ids)
 
    # 加入校準雜訊以擾亂浮水印訊號
    noise = torch.randn_like(embeddings) * noise_scale
    perturbed_embeddings = embeddings + noise
 
    # 自擾動後之 embedding 重新解碼
    with torch.no_grad():
        outputs = model(inputs_embeds=perturbed_embeddings)
        new_tokens = outputs.logits.argmax(dim=-1)
 
    return tokenizer.decode(new_tokens[0], skip_special_tokens=True)