From c6739526c60fac2df39f0440f8c54549dfa8bb9d Mon Sep 17 00:00:00 2001
From: Karol Damaszke <karol.damaszke@intel.com>
Date: Tue, 9 Apr 2024 11:29:21 +0200
Subject: [PATCH] Fix test_watermark (#124)

Co-authored-by: Karol Damaszke <kdamaszke@habana.ai>
---
 server/tests/utils/test_tokens.py    |  5 +-
 server/tests/utils/test_watermark.py | 85 +++++++++++++++++++++-------
 2 files changed, 66 insertions(+), 24 deletions(-)

diff --git a/server/tests/utils/test_tokens.py b/server/tests/utils/test_tokens.py
index 7f78c982..39446fd5 100644
--- a/server/tests/utils/test_tokens.py
+++ b/server/tests/utils/test_tokens.py
@@ -1,6 +1,9 @@
 # Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
 
+import pytest
 import torch
+from transformers import AutoTokenizer
+
 from text_generation_server.utils.tokens import (
     StopSequenceCriteria,
     StoppingCriteria,
@@ -8,10 +11,8 @@ from text_generation_server.utils.tokens import (
     batch_top_tokens,
     make_tokenizer_optional,
 )
-from transformers import AutoTokenizer
 
 
-import pytest
 @pytest.fixture
 def skip_tokenizer_env_var():
     import os
diff --git a/server/tests/utils/test_watermark.py b/server/tests/utils/test_watermark.py
index 5d0aaf05..c7c83e8c 100644
--- a/server/tests/utils/test_watermark.py
+++ b/server/tests/utils/test_watermark.py
@@ -1,7 +1,8 @@
-# test_watermark_logits_processor.py
+# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.
 
 import os
 import numpy as np
+import pytest
 import torch
 from text_generation_server.utils.watermark import WatermarkLogitsProcessor
 
@@ -10,35 +11,71 @@ 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)
+@pytest.fixture
+def hpu_device():
+    return torch.device("hpu")
+
+
+@pytest.fixture
+def input_ids_list():
+    return [101, 2036, 3731, 102, 2003, 103]
+
+
+@pytest.fixture
+def input_ids_tensor(hpu_device):
+    return torch.tensor(
+        [[101, 2036, 3731, 102, 2003, 103]],
+        dtype=torch.int64,
+        device=hpu_device
+    )
+
+
+@pytest.fixture
+def scores(hpu_device):
+    return torch.tensor(
+        [[0.5, 0.3, 0.2, 0.8], [0.1, 0.2, 0.7, 0.9]],
+        device=hpu_device
+    )
+
+
+def test_seed_rng(input_ids_list, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    processor._seed_rng(input_ids_list)
     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"))
+def test_seed_rng_tensor(input_ids_tensor, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    processor._seed_rng(input_ids_tensor)
+    assert isinstance(processor.rng, torch.Generator)
+
+
+def test_get_greenlist_ids(input_ids_list, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    result = processor._get_greenlist_ids(input_ids_list, 10, hpu_device)
     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])
+def test_get_greenlist_ids_tensor(input_ids_tensor, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    result = processor._get_greenlist_ids(input_ids_tensor, 10, hpu_device)
+    assert max(result) <= 10
+    assert len(result) == int(10 * 0.5)
+
+
+def test_calc_greenlist_mask(scores, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    greenlist_token_ids = torch.tensor([2, 3], device=hpu_device)
     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]])
+def test_bias_greenlist_logits(scores, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
     green_tokens_mask = torch.tensor(
-        [[False, False, True, True], [False, False, False, True]]
+        [[False, False, True, True], [False, False, False, True]], device=hpu_device
     )
     greenlist_bias = 2.0
     result = processor._bias_greenlist_logits(scores, green_tokens_mask, greenlist_bias)
@@ -46,9 +83,13 @@ def test_bias_greenlist_logits():
     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)
+def test_call(input_ids_list, scores, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    result = processor(input_ids_list, scores)
+    assert result.shape == scores.shape
+
+
+def test_call_tensor(input_ids_tensor, scores, hpu_device):
+    processor = WatermarkLogitsProcessor(device=hpu_device)
+    result = processor(input_ids_tensor, scores)
     assert result.shape == scores.shape