added test for generate(strip_prompt=True)

keras-team · Oct 9, 2024 · b098a0b · b098a0b
1 parent 16f6109
commit b098a0b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/keras_hub/src/models/llama3/llama3_causal_lm_test.py b/keras_hub/src/models/llama3/llama3_causal_lm_test.py
@@ -69,6 +69,12 @@ def test_generate(self):
             prompt_ids["padding_mask"][:, :5],
         )
 
+    def test_generate_strip_prompt(self):
+        causal_lm = Llama3CausalLM(**self.init_kwargs)
+        prompt = " airplane at airport"
+        output = causal_lm.generate(prompt, strip_prompt=True)
+        self.assertFalse(prompt in output)
+
     def test_early_stopping(self):
         causal_lm = Llama3CausalLM(**self.init_kwargs)
         call_with_cache = causal_lm.call_with_cache