[eagle2] fix end check when target model verify (#2723)

test/srt/test_eagle_infer.py

@@ -1,5 +1,7 @@
 import unittest
 
+from transformers import AutoConfig, AutoTokenizer
+
 import sglang as sgl
 
 
@@ -34,6 +36,33 @@ class TestEAGLEEngine(unittest.TestCase):
         print(out2)
         self.assertEqual(out1, out2)
 
+    def test_eagle_end_check(self):
+        prompt = "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nToday is a sunny day and I like [/INST]"
+        target_model_path = "meta-llama/Llama-2-7b-chat-hf"
+        tokenizer = AutoTokenizer.from_pretrained(target_model_path)
+        speculative_draft_model_path = "lmzheng/sglang-EAGLE-llama2-chat-7B"
+
+        sampling_params = {
+            "temperature": 0,
+            "max_new_tokens": 1024,
+            "skip_special_tokens": False,
+        }
+
+        engine = sgl.Engine(
+            model_path=target_model_path,
+            speculative_draft_model_path=speculative_draft_model_path,
+            speculative_algorithm="EAGLE",
+            speculative_num_steps=3,
+            speculative_eagle_topk=4,
+            speculative_num_draft_tokens=16,
+        )
+        out1 = engine.generate(prompt, sampling_params)["text"]
+        engine.shutdown()
+        print("==== Answer 1 ====")
+        print(repr(out1))
+        tokens = tokenizer.encode(out1, truncation=False)
+        assert tokenizer.eos_token_id not in tokens
+
 
 if __name__ == "__main__":
     unittest.main()