7ee8a6dd01
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Signed-off-by: Simon Veitner <sveitner@redhat.com> Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>
23 lines
778 B
Python
23 lines
778 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from vllm.config.watermarking import WatermarkConfig
|
|
from vllm.v1.watermarking.gumbel import (
|
|
DualKeyGumbelWatermarker,
|
|
GumbelWatermarker,
|
|
)
|
|
from vllm.v1.watermarking.watermarker import Watermarker
|
|
|
|
|
|
def create_watermarker(config: WatermarkConfig) -> Watermarker:
|
|
if config.algorithm == "gumbel":
|
|
return GumbelWatermarker(config.key, config.context_width, config.prf)
|
|
if config.algorithm == "dual_key_gumbel":
|
|
return DualKeyGumbelWatermarker(
|
|
config.key,
|
|
config.context_width,
|
|
config.prf,
|
|
config.alpha,
|
|
)
|
|
raise ValueError(f"Unknown watermarking algorithm: {config.algorithm}")
|