ea40bb9e90
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
20 lines
522 B
Python
20 lines
522 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from vllm.config.watermarking import WatermarkPRFName
|
|
from vllm.v1.watermarking.prfs.base import WatermarkPRF
|
|
from vllm.v1.watermarking.prfs.philox import PhiloxPRF
|
|
|
|
|
|
def create_prf(name: WatermarkPRFName, key: int) -> WatermarkPRF:
|
|
if name == "philox":
|
|
return PhiloxPRF(key)
|
|
raise ValueError(f"Unknown watermark PRF: {name}")
|
|
|
|
|
|
__all__ = [
|
|
"PhiloxPRF",
|
|
"WatermarkPRF",
|
|
"create_prf",
|
|
]
|