| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm.config.watermarking import WatermarkPRFName |
| 5 | from vllm.v1.watermarking.prfs.base import WatermarkPRF |
| 6 | from vllm.v1.watermarking.prfs.philox import PhiloxPRF |
| 7 | |
| 8 | |
| 9 | def create_prf(name: WatermarkPRFName, key: int) -> WatermarkPRF: |
| 10 | if name == "philox": |
| 11 | return PhiloxPRF(key) |
| 12 | raise ValueError(f"Unknown watermark PRF: {name}") |
| 13 | |
| 14 | |
| 15 | __all__ = [ |
| 16 | "PhiloxPRF", |
| 17 | "WatermarkPRF", |
| 18 | "create_prf", |
| 19 | ] |