| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | _UINT32_MASK = 2**32 - 1 |
| 9 | |
| 10 | |
| 11 | class WatermarkPRF(ABC): |
| 12 | @abstractmethod |
| 13 | def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: |
| 14 | raise NotImplementedError |
| 15 | |
| 16 | |
| 17 | def uint32_to_uniform(values: torch.Tensor) -> torch.Tensor: |
| 18 | mantissas = (values & _UINT32_MASK) >> 8 |
| 19 | return ((mantissas.to(torch.float64) + 1.0) / (2**24 + 1)).to(torch.float32) |