| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from dataclasses import dataclass |
| 6 | from typing import TypeAlias |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | Context: TypeAlias = tuple[int, ...] |
| 11 | |
| 12 | |
| 13 | @dataclass(frozen=True) |
| 14 | class WatermarkDetection: |
| 15 | score: float |
| 16 | p_value: float |
| 17 | num_scored_tokens: int |
| 18 | is_watermarked: bool |
| 19 | |
| 20 | |
| 21 | class WatermarkDetector(ABC): |
| 22 | def __init__( |
| 23 | self, |
| 24 | context_width: int, |
| 25 | p_value_threshold: float, |
| 26 | deduplicate_contexts: bool = False, |
| 27 | ) -> None: |
| 28 | if context_width < 1: |
| 29 | raise ValueError("context_width must be positive") |
| 30 | if not 0 < p_value_threshold < 1: |
| 31 | raise ValueError("p_value_threshold must be between 0 and 1") |
| 32 | self.context_width = context_width |
| 33 | self.p_value_threshold = p_value_threshold |
| 34 | self.deduplicate_contexts = deduplicate_contexts |
| 35 | |
| 36 | def detect(self, token_ids: list[int]) -> WatermarkDetection: |
| 37 | contexts, targets = self._prepare_inputs(token_ids) |
| 38 | num_scored_tokens = len(targets) |
| 39 | if num_scored_tokens == 0: |
| 40 | return WatermarkDetection(0.0, 1.0, 0, False) |
| 41 | |
| 42 | token_scores = self._score_tokens(contexts, targets) |
| 43 | score = self._aggregate_scores(token_scores) |
| 44 | p_value = self._get_p_value(score, num_scored_tokens) |
| 45 | return WatermarkDetection( |
| 46 | score=score, |
| 47 | p_value=p_value, |
| 48 | num_scored_tokens=num_scored_tokens, |
| 49 | is_watermarked=p_value <= self.p_value_threshold, |
| 50 | ) |
| 51 | |
| 52 | def _prepare_inputs( |
| 53 | self, token_ids: list[int] |
| 54 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 55 | contexts: list[Context] = [] |
| 56 | targets = [] |
| 57 | seen_contexts: set[Context] = set() |
| 58 | prefix = [-1] * self.context_width |
| 59 | for token_id in token_ids: |
| 60 | context = tuple(prefix[-self.context_width :]) |
| 61 | # Reusing a context reuses one keyed random vector, so the resulting |
| 62 | # token scores are not independent evidence. |
| 63 | if not self.deduplicate_contexts or not _is_repeated_context( |
| 64 | context, seen_contexts |
| 65 | ): |
| 66 | contexts.append(context) |
| 67 | targets.append(token_id) |
| 68 | prefix.append(token_id) |
| 69 | |
| 70 | if not contexts: |
| 71 | return ( |
| 72 | torch.empty((0, self.context_width), dtype=torch.int64), |
| 73 | torch.empty(0, dtype=torch.int64), |
| 74 | ) |
| 75 | |
| 76 | return ( |
| 77 | torch.tensor(contexts, dtype=torch.int64).reshape(-1, self.context_width), |
| 78 | torch.tensor(targets, dtype=torch.int64), |
| 79 | ) |
| 80 | |
| 81 | @abstractmethod |
| 82 | def _score_tokens( |
| 83 | self, contexts: torch.Tensor, targets: torch.Tensor |
| 84 | ) -> torch.Tensor: |
| 85 | raise NotImplementedError |
| 86 | |
| 87 | @abstractmethod |
| 88 | def _aggregate_scores(self, token_scores: torch.Tensor) -> float: |
| 89 | raise NotImplementedError |
| 90 | |
| 91 | @abstractmethod |
| 92 | def _get_p_value(self, score: float, num_scored_tokens: int) -> float: |
| 93 | raise NotImplementedError |
| 94 | |
| 95 | |
| 96 | def _is_repeated_context(context: Context, seen_contexts: set[Context]) -> bool: |
| 97 | if context in seen_contexts: |
| 98 | return True |
| 99 | seen_contexts.add(context) |
| 100 | return False |