55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
# test_watermark_logits_processor.py
|
|
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
from text_generation_server.utils.watermark import WatermarkLogitsProcessor
|
|
|
|
|
|
GAMMA = os.getenv("WATERMARK_GAMMA", 0.5)
|
|
DELTA = os.getenv("WATERMARK_DELTA", 2.0)
|
|
|
|
|
|
def test_seed_rng():
|
|
input_ids = [101, 2036, 3731, 102, 2003, 103]
|
|
processor = WatermarkLogitsProcessor()
|
|
processor._seed_rng(input_ids)
|
|
assert isinstance(processor.rng, torch.Generator)
|
|
|
|
|
|
def test_get_greenlist_ids():
|
|
input_ids = [101, 2036, 3731, 102, 2003, 103]
|
|
processor = WatermarkLogitsProcessor()
|
|
result = processor._get_greenlist_ids(input_ids, 10, torch.device("cpu"))
|
|
assert max(result) <= 10
|
|
assert len(result) == int(10 * 0.5)
|
|
|
|
|
|
def test_calc_greenlist_mask():
|
|
processor = WatermarkLogitsProcessor()
|
|
scores = torch.tensor([[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]])
|
|
greenlist_token_ids = torch.tensor([2, 3])
|
|
result = processor._calc_greenlist_mask(scores, greenlist_token_ids)
|
|
assert result.tolist() == [[False, False, False, False], [False, False, True, True]]
|
|
assert result.shape == scores.shape
|
|
|
|
|
|
def test_bias_greenlist_logits():
|
|
processor = WatermarkLogitsProcessor()
|
|
scores = torch.tensor([[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]])
|
|
green_tokens_mask = torch.tensor(
|
|
[[False, False, True, True], [False, False, False, True]]
|
|
)
|
|
greenlist_bias = 2.0
|
|
result = processor._bias_greenlist_logits(scores, green_tokens_mask, greenlist_bias)
|
|
assert np.allclose(result.tolist(), [[0.5, 0.3, 2.2, 2.8], [0.1, 0.2, 0.7, 2.9]])
|
|
assert result.shape == scores.shape
|
|
|
|
|
|
def test_call():
|
|
input_ids = [101, 2036, 3731, 102, 2003, 103]
|
|
processor = WatermarkLogitsProcessor()
|
|
scores = torch.tensor([[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]])
|
|
result = processor(input_ids, scores)
|
|
assert result.shape == scores.shape
|