| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm.config.watermarking import WatermarkConfig |
| 5 | from vllm.v1.watermarking.gumbel import ( |
| 6 | DualKeyGumbelWatermarker, |
| 7 | GumbelWatermarker, |
| 8 | ) |
| 9 | from vllm.v1.watermarking.watermarker import Watermarker |
| 10 | |
| 11 | |
| 12 | def create_watermarker(config: WatermarkConfig) -> Watermarker: |
| 13 | if config.algorithm == "gumbel": |
| 14 | return GumbelWatermarker(config.key, config.context_width, config.prf) |
| 15 | if config.algorithm == "dual_key_gumbel": |
| 16 | return DualKeyGumbelWatermarker( |
| 17 | config.key, |
| 18 | config.context_width, |
| 19 | config.prf, |
| 20 | config.alpha, |
| 21 | ) |
| 22 | raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}") |