Fix tests with do_sample=True (#30)

* fix(tests): update expected result when do_sample=True * chore: try to set seed on xla to stabilize tests
huggingface · Apr 29, 2024 · f9a6593 · f9a6593
1 parent 5e5a1d7
commit f9a6593
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/optimum/tpu/generation/token_selector.py b/optimum/tpu/generation/token_selector.py
@@ -3,6 +3,7 @@
 from typing import Optional
 
 import torch
+import torch_xla.core.xla_model as xm
 from transformers.generation import (
     GenerationConfig,
     GenerationMixin,
@@ -52,6 +53,7 @@ def __init__(
         self.eos_token_id = eos_token_id
         self.pad_token_id = pad_token_id
         self.logits_warper = logits_warper
+        xm.set_rng_state(seed)
         self.generator = torch.Generator()
         self.generator.manual_seed(seed)
 

diff --git a/text-generation-inference/tests/test_gpt2.py b/text-generation-inference/tests/test_gpt2.py
@@ -72,8 +72,8 @@ def create_request(
         ],
         [
             "It was a bright cold day in April, and the clocks were striking thirteen.",
-            775,
-            " We",
+            1439,
+            " All",
             True,
         ],
     ],
@@ -113,7 +113,7 @@ def test_prefill(input_text, token_id, token_text, do_sample, batch_size, model_
         [
             "It was a bright cold day in April, and the clocks were striking thirteen.",
             20,
-            " We sat at the front door and watched the clock on a box of yellow paper and found it almost",
+            " All day the sun had set, as was well-known. The first thing I noticed when I",
             True,
         ],
     ],
-Original file line number
+Diff line change
@@ Expand Up / @@ -72,8 +72,8 @@ def create_request( @@
             ],
             [
                 "It was a bright cold day in April, and the clocks were striking thirteen.",
-,
-                " We",
+,
+                " All",
                 True,
             ],
         ],
@@ Expand Down Expand Up @@
             [
                 "It was a bright cold day in April, and the clocks were striking thirteen.",
 ,
-                " We sat at the front door and watched the clock on a box of yellow paper and found it almost",
+                " All day the sun had set, as was well-known. The first thing I noticed when I",
                 True,
             ],
         ],
@@ Expand Down @@