fix bad_words test

Author: Sunny-bot1

test/ci_use/EB_Lite/test_EB_Lite_serving.py

@@ -784,158 +784,103 @@ def test_non_streaming_chat_with_bad_words(openai_client, capsys):
     """
     Test bad_words option in non-streaming chat functionality with the local service
     """
+    base_path = os.getenv("MODEL_PATH")
+    if base_path:
+        model_path = os.path.join(base_path, "ernie-4_5-21b-a3b-bf16-paddle")
+    else:
+        model_path = "./ernie-4_5-21b-a3b-bf16-paddle"
     response_0 = openai_client.chat.completions.create(
         model="default",
         messages=[{"role": "user", "content": "Hello, how are you?"}],
         temperature=1,
         top_p=0.0,
-        max_tokens=10,
+        max_tokens=20,
         stream=False,
+        extra_body={"return_token_ids": True},
     )
-    output_0 = []
+
     assert hasattr(response_0, "choices")
     assert len(response_0.choices) > 0
     assert hasattr(response_0.choices[0], "message")
-    assert hasattr(response_0.choices[0].message, "content")
+    assert hasattr(response_0.choices[0].message, "completion_token_ids")
+    assert isinstance(response_0.choices[0].message.completion_token_ids, list)
 
-    text_split = response_0.choices[0].message.content.split(" ")
-    for text in text_split:
-        output_0.append(text)
+    from fastdeploy.input.ernie_tokenizer import ErnieBotTokenizer
+
+    tokenizer = ErnieBotTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    output_tokens_0 = []
+    output_ids_0 = []
+    for ids in response_0.choices[0].message.completion_token_ids:
+        output_tokens_0.append(tokenizer.decode(ids))
+        output_ids_0.append(ids)
 
     # add bad words
+    bad_tokens = output_tokens_0[6:10]
+    bad_token_ids = output_ids_0[6:10]
     response_1 = openai_client.chat.completions.create(
         model="default",
         messages=[{"role": "user", "content": "Hello, how are you?"}],
         temperature=1,
         top_p=0.0,
-        max_tokens=10,
-        extra_body={"bad_words": output_0[-5:]},
+        max_tokens=20,
+        extra_body={"bad_words": bad_tokens[-10:], "return_token_ids": True},
         stream=False,
     )
-    output_1 = []
     assert hasattr(response_1, "choices")
     assert len(response_1.choices) > 0
     assert hasattr(response_1.choices[0], "message")
-    assert hasattr(response_1.choices[0].message, "content")
-    text_split = response_1.choices[0].message.content.split(" ")
-    for text in text_split:
-        output_1.append(text)
-    assert output_0 not in output_1
-
-
-def test_streaming_chat_with_bad_words(openai_client, capsys):
-    """
-    Test bad_words option in streaming chat functionality with the local service
-    """
-    response_0 = openai_client.chat.completions.create(
-        model="default",
-        messages=[{"role": "user", "content": "Hello, how are you?"}],
-        temperature=1,
-        top_p=0.0,
-        max_tokens=10,
-        stream=True,
-    )
-    output_0 = []
-    for chunk in response_0:
-        assert hasattr(chunk, "choices")
-        assert len(chunk.choices) > 0
-        assert hasattr(chunk.choices[0], "delta")
-        assert hasattr(chunk.choices[0].delta, "content")
-        output_0.append(chunk.choices[0].delta.content)
-
-    # add bad words
-    response_1 = openai_client.chat.completions.create(
-        model="default",
-        messages=[{"role": "user", "content": "Hello, how are you?"}],
-        temperature=1,
-        top_p=0.0,
-        max_tokens=10,
-        extra_body={"bad_words": output_0[-5:]},
-        stream=True,
-    )
-    output_1 = []
-    for chunk in response_1:
-        assert hasattr(chunk, "choices")
-        assert len(chunk.choices) > 0
-        assert hasattr(chunk.choices[0], "delta")
-        assert hasattr(chunk.choices[0].delta, "content")
-        output_1.append(chunk.choices[0].delta.content)
-    assert output_0 not in output_1
+    assert hasattr(response_1.choices[0].message, "completion_token_ids")
+    assert isinstance(response_1.choices[0].message.completion_token_ids, list)
+    assert not any(ids in response_1.choices[0].message.completion_token_ids for ids in bad_token_ids)
 
 
 def test_non_streaming_completion_with_bad_words(openai_client, capsys):
     """
     Test bad_words option in non-streaming completion functionality with the local service
     """
+    base_path = os.getenv("MODEL_PATH")
+    if base_path:
+        model_path = os.path.join(base_path, "ernie-4_5-21b-a3b-bf16-paddle")
+    else:
+        model_path = "./ernie-4_5-21b-a3b-bf16-paddle"
+
     response_0 = openai_client.completions.create(
         model="default",
         prompt="Hello, how are you?",
         temperature=1,
         top_p=0.0,
-        max_tokens=10,
+        max_tokens=20,
         stream=False,
+        extra_body={"return_token_ids": True},
     )
-    output_0 = []
     assert hasattr(response_0, "choices")
     assert len(response_0.choices) > 0
-    assert hasattr(response_0.choices[0], "text")
-    text_split = response_0.choices[0].text.split(" ")
-    for text in text_split:
-        output_0.append(text)
+    assert hasattr(response_0.choices[0], "completion_token_ids")
+    assert isinstance(response_0.choices[0].completion_token_ids, list)
+
+    from fastdeploy.input.ernie_tokenizer import ErnieBotTokenizer
+
+    tokenizer = ErnieBotTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    output_tokens_0 = []
+    output_ids_0 = []
+    for ids in response_0.choices[0].completion_token_ids:
+        output_tokens_0.append(tokenizer.decode(ids))
+        output_ids_0.append(ids)
 
     # add bad words
+    bad_tokens = output_tokens_0[6:10]
+    bad_token_ids = output_ids_0[6:10]
     response_1 = openai_client.completions.create(
         model="default",
         prompt="Hello, how are you?",
         temperature=1,
         top_p=0.0,
-        max_tokens=10,
-        extra_body={"bad_words": output_0[-5:]},
+        max_tokens=20,
+        extra_body={"bad_words": bad_tokens, "return_token_ids": True},
         stream=False,
     )
-    output_1 = []
     assert hasattr(response_1, "choices")
     assert len(response_1.choices) > 0
-    assert hasattr(response_1.choices[0], "text")
-    text_split = response_1.choices[0].text.split(" ")
-    for text in text_split:
-        output_1.append(text)
-    assert output_0 not in output_1
-
-
-def test_streaming_completion_with_bad_words(openai_client, capsys):
-    """
-    Test bad_words option in streaming completion functionality with the local service
-    """
-    response_0 = openai_client.completions.create(
-        model="default",
-        prompt="Hello, how are you?",
-        temperature=1,
-        top_p=0.0,
-        max_tokens=10,
-        stream=True,
-    )
-    output_0 = []
-    for chunk in response_0:
-        assert hasattr(chunk, "choices")
-        assert len(chunk.choices) > 0
-        assert hasattr(chunk.choices[0], "text")
-        output_0.append(chunk.choices[0].text)
-
-    # add bad words
-    response_1 = openai_client.completions.create(
-        model="default",
-        prompt="Hello, how are you?",
-        temperature=1,
-        top_p=0.0,
-        max_tokens=10,
-        extra_body={"bad_words": output_0[-5:]},
-        stream=True,
-    )
-    output_1 = []
-    for chunk in response_1:
-        assert hasattr(chunk, "choices")
-        assert len(chunk.choices) > 0
-        assert hasattr(chunk.choices[0], "text")
-        output_1.append(chunk.choices[0].text)
-    assert output_0 not in output_1
+    assert hasattr(response_0.choices[0], "completion_token_ids")
+    assert isinstance(response_0.choices[0].completion_token_ids, list)
+    assert not any(ids in response_1.choices[0].completion_token_ids for ids in bad_token_ids)