AI Watermarking

"""
Phát hiện watermark Kirchenbauer bằng z-score.
Yêu cầu: tokenizer của mô hình gốc, gamma đã biết.
"""
import hashlib
import numpy as np
from typing import List, Dict

def _green_list_for_context(
    prev_token_id: int,
    vocab_size: int,
    gamma: float = 0.5,
    secret_key: bytes = b"vi-edu-demo-key",
) -> set[int]:
    """Sinh green list từ hash(khoá + prev_token).
    Trả về tập các token id thuộc green list.
    """
    payload = secret_key + prev_token_id.to_bytes(8, "little")
    seed_bytes = hashlib.sha256(payload).digest()[:8]
    seed = int.from_bytes(seed_bytes, "little")
    rng = np.random.default_rng(seed)

    green_size = int(vocab_size * gamma)
    green_ids = rng.choice(vocab_size, size=green_size, replace=False)
    return set(int(x) for x in green_ids)


def detect_watermark(
    token_ids: List[int],
    vocab_size: int,
    gamma: float = 0.5,
    secret_key: bytes = b"vi-edu-demo-key",
    z_threshold: float = 4.0,
) -> Dict[str, float | bool | str]:
    """Phát hiện watermark trong một chuỗi token."""
    if len(token_ids) < 2:
        raise ValueError("Cần ít nhất 2 token để phát hiện watermark.")

    green_count = 0
    total = len(token_ids) - 1  # Bỏ token đầu (không có prev)

    for i in range(1, len(token_ids)):
        prev = token_ids[i - 1]
        green_list = _green_list_for_context(
            prev, vocab_size, gamma, secret_key,
        )
        if token_ids[i] in green_list:
            green_count += 1

    # z-score test
    expected = total * gamma
    variance = total * gamma * (1 - gamma)
    z = (green_count - expected) / np.sqrt(variance)

    # p-value xấp xỉ (1-tail)
    from scipy.stats import norm
    p_value = 1.0 - norm.cdf(z)

    confidence = (
        "cao" if z > 6.0
        else "vừa" if z > z_threshold
        else "thấp"
    )

    return {
        "z_score": float(z),
        "p_value": float(p_value),
        "green_ratio": green_count / total,
        "is_ai": bool(z > z_threshold),
        "confidence": confidence,
        "total_tokens": total,
    }


# ─────────────────────────────────────────────────────────
# Ví dụ sử dụng với tokenizer thực tế
# ─────────────────────────────────────────────────────────
if __name__ == "__main__":
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained("vinai/PhoGPT-7B5")
    text = open("suspicious_essay.txt", encoding="utf-8").read()
    ids = tok.encode(text)

    result = detect_watermark(
        token_ids=ids,
        vocab_size=tok.vocab_size,
        gamma=0.5,
    )

    print(f"Tổng tokens:     {result['total_tokens']}")
    print(f"Tỷ lệ green:     {result['green_ratio']:.3f}")
    print(f"z-score:         {result['z_score']:.2f}")
    print(f"p-value:         {result['p_value']:.2e}")
    print(f"Phán đoán:       {'AI' if result['is_ai'] else 'Người/mơ hồ'}")
    print(f"Độ tin cậy:      {result['confidence']}")

"""
Sinh văn bản có watermark bằng logit processor.
Tương thích với Hugging Face Transformers.
"""
import torch
from transformers import LogitsProcessor


class WatermarkLogitsProcessor(LogitsProcessor):
    """Cộng delta vào logits của green tokens trước softmax.

    Triển khai theo Kirchenbauer et al., 2023.
    """

    def __init__(
        self,
        vocab_size: int,
        gamma: float = 0.5,
        delta: float = 2.0,
        secret_key: bytes = b"vi-edu-demo-key",
    ):
        self.vocab_size = vocab_size
        self.gamma = gamma
        self.delta = delta
        self.secret_key = secret_key
        self._cache: dict[int, torch.Tensor] = {}

    def _green_mask(self, prev_token: int, device: torch.device) -> torch.Tensor:
        """Trả về mask nhị phân (vocab_size,) — 1 cho green, 0 cho red."""
        if prev_token in self._cache:
            return self._cache[prev_token].to(device)

        import hashlib, numpy as np
        payload = self.secret_key + int(prev_token).to_bytes(8, "little")
        seed = int.from_bytes(
            hashlib.sha256(payload).digest()[:8], "little",
        )
        rng = np.random.default_rng(seed)
        size = int(self.vocab_size * self.gamma)
        green = rng.choice(self.vocab_size, size=size, replace=False)

        mask = torch.zeros(self.vocab_size)
        mask[green] = 1.0
        self._cache[prev_token] = mask
        return mask.to(device)

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
    ) -> torch.FloatTensor:
        # input_ids: (batch, seq_len) — seq_len có thể khác nhau giữa batch
        # scores: (batch, vocab_size) — logits cho token tiếp theo
        for i in range(input_ids.shape[0]):
            prev = int(input_ids[i, -1].item())
            mask = self._green_mask(prev, scores.device)
            scores[i] = scores[i] + self.delta * mask
        return scores


# ─────────────────────────────────────────────────────────
# Ví dụ: watermark hoá PhoGPT tiếng Việt
# ─────────────────────────────────────────────────────────
if __name__ == "__main__":
    from transformers import AutoModelForCausalLM, AutoTokenizer

    name = "vinai/PhoGPT-7B5-Instruct"
    tok = AutoTokenizer.from_pretrained(name)
    model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=torch.float16)

    processor = WatermarkLogitsProcessor(
        vocab_size=tok.vocab_size, gamma=0.5, delta=2.0,
    )

    inputs = tok("Giới thiệu về vịnh Hạ Long:", return_tensors="pt")
    out = model.generate(
        **inputs,
        max_new_tokens=200,
        logits_processor=[processor],
        do_sample=True,
        top_p=0.9,
    )
    print(tok.decode(out[0], skip_special_tokens=True))