Fix tests with the intended "auto" default (keras-team#1559)

keras_nlp/models/bart/bart_seq_2_seq_lm_test.py

@@ -87,7 +87,7 @@ def test_generate(self):
         # Int tensor input.
         seq_2_seq_lm.preprocessor = None
         preprocessed_batch = self.preprocessor.generate_preprocess(inputs)
-        outputs = seq_2_seq_lm.generate(preprocessed_batch)
+        outputs = seq_2_seq_lm.generate(preprocessed_batch, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["decoder_token_ids"][:, :5],

keras_nlp/models/bloom/bloom_causal_lm_test.py

@@ -78,7 +78,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :4],
@@ -103,7 +103,7 @@ def test_generate_with_bfloat16(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :4],
@@ -128,7 +128,7 @@ def test_generate_with_mixed_float16(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :4],

keras_nlp/models/gemma/gemma_causal_lm_test.py

@@ -73,7 +73,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :4],
@@ -96,7 +96,7 @@ def test_generate_with_bfloat16(self):
             # Int tensor input.
             prompt_ids = self.preprocessor.generate_preprocess([prompt])
             causal_lm.preprocessor = None
-            outputs = causal_lm.generate(prompt_ids)
+            outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
             # Assert prompt is in output in token id space.
             self.assertAllEqual(
                 outputs["token_ids"][:, :4],

keras_nlp/models/gpt2/gpt2_causal_lm_test.py

@@ -70,7 +70,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :5],

keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm_test.py

@@ -70,7 +70,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :5],

keras_nlp/models/llama/llama_causal_lm.py

@@ -187,9 +187,8 @@ def next(prompt, cache, index):
         if stop_token_ids is not None:
             # Build a mask of stop token locations not in the original
             # prompt (not in locations where `padding_mask` is True).
-            end_locations = ops.logical_and(
-                any_equal(token_ids, stop_token_ids),
-                ops.logical_not(padding_mask),
+            end_locations = any_equal(
+                token_ids, stop_token_ids, ops.logical_not(padding_mask)
             )
             end_locations = ops.cast(end_locations, "int32")
             # Use cumsum to get ones in all locations after end_locations.

keras_nlp/models/llama/llama_causal_lm_test.py

@@ -70,7 +70,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :5],

keras_nlp/models/mistral/mistral_causal_lm_test.py

@@ -70,7 +70,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :5],

keras_nlp/models/opt/opt_causal_lm_test.py

@@ -69,7 +69,7 @@ def test_generate(self):
         # Int tensor input.
         prompt_ids = self.preprocessor.generate_preprocess([prompt])
         causal_lm.preprocessor = None
-        outputs = causal_lm.generate(prompt_ids)
+        outputs = causal_lm.generate(prompt_ids, stop_token_ids=None)
         # Assert prompt is in output in token id space.
         self.assertAllEqual(
             outputs["token_ids"][:, :5],