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
542 B
Python
20 lines
542 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
import torch
|
|
|
|
_UINT32_MASK = 2**32 - 1
|
|
|
|
|
|
class WatermarkPRF(ABC):
|
|
@abstractmethod
|
|
def uniform(self, contexts: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
|
|
raise NotImplementedError
|
|
|
|
|
|
def uint32_to_uniform(values: torch.Tensor) -> torch.Tensor:
|
|
mantissas = (values & _UINT32_MASK) >> 8
|
|
return ((mantissas.to(torch.float64) + 1.0) / (2**24 + 1)).to(torch.float32)
|