"""Streamer tests in MLC LLM. Please specify the local path to llama2 tokenizer via environment variable before running this test. The recommended way to run the tests is to use the following command: MLC_LLAMA_TOKENIZER_PATH="path/to/llama/tokenizer" \ pytest -vv tests/python/support/test_text_streamer_stop_handler.py Here "MLC_LLAMA_TOKENIZER_PATH" can be chosen from - a llama2 weight directory (e.g., "path/to/Llama-2-7b-chat-hf"), - a sentencepiece llama2 tokenizer path (e.g., "path/to/Llama-2-7b-chat-hf/tokenizer.model"). To directly run the Python file (a.k.a., not using pytest), you also need to specify the tokenizer path via environment variable. """ import time from typing import List, Tuple # noqa: UP035 import pytest from mlc_llm.testing import require_test_tokenizers from mlc_llm.tokenizers import StopStrHandler, TextStreamer, Tokenizer # test category "unittest" pytestmark = [pytest.mark.unittest] # fmt: off para_input_tokens = [18585, 29892, 1244, 29915, 29879, 263, 3273, 14880, 1048, 953, 29877, 2397, 29892, 988, 1269, 1734, 338, 5643, 491, 385, 953, 29877, 2397, 29901, 13, 13, 29950, 1032, 727, 29991, 29871, 243, 162, 148, 142, 306, 29915, 29885, 1244, 304, 1371, 1234, 738, 5155, 366, 505, 1048, 953, 29877, 2397, 29871, 243, 162, 167, 151, 29889, 7440, 366, 1073, 393, 953, 29877, 2397, 508, 367, 1304, 304, 27769, 23023, 1080, 322, 21737, 297, 263, 2090, 322, 1708, 1319, 982, 29973, 29871, 243, 162, 155, 135, 2688, 508, 884, 367, 1304, 304, 788, 263, 6023, 310, 2022, 2877, 304, 596, 7191, 322, 11803, 29889, 29871, 243, 162, 149, 152, 1126, 29892, 1258, 366, 1073, 393, 727, 526, 1584, 953, 29877, 2397, 8090, 322, 14188, 366, 508, 1708, 29973, 29871, 243, 162, 145, 177, 243, 162, 148, 131, 1105, 29892, 748, 14432, 322, 679, 907, 1230, 411, 953, 29877, 2397, 29991, 29871, 243, 162, 149, 168, 243, 162, 145, 171] DECODED_PARAGRAPH = ( "Sure, here's a short paragraph about emoji, " "where each word is followed by an emoji:\n\n" "Hey there! 👋 I'm here to help answer any questions you have about emoji 🤔. " "Did you know that emoji can be used to convey emotions and feelings in a " "fun and playful way? 😄 " "They can also be used to add a touch of personality to your messages and posts. 💕 " "And, did you know that there are even emoji games and activities you can play? 🎮👀 " "So, go ahead and get creative with emoji! 💥🎨" ) # fmt: on @require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") def test_text_streamer(llama_tokenizer_path: str): text_streamer = TextStreamer(Tokenizer(llama_tokenizer_path)) total_text = "" for token in para_input_tokens: total_text += text_streamer.put([token]) total_text += text_streamer.finish() assert total_text == DECODED_PARAGRAPH def stop_handler_process_tokens( stop_handler: StopStrHandler, tokens: List[int], # noqa: UP006 tokenizer: Tokenizer, ) -> str: returned_tokens = [] for token in tokens: returned_tokens += stop_handler.put(token) if stop_handler.stop_triggered: break if not stop_handler.stop_triggered: returned_tokens += stop_handler.finish() return tokenizer.decode(returned_tokens) @require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") def test_stop_str_handler_stop(llama_tokenizer_path: str): stop_strs = [" 🤔"] tokenizer = Tokenizer(llama_tokenizer_path) stop_handler = StopStrHandler(stop_strs, tokenizer) total_text = stop_handler_process_tokens(stop_handler, para_input_tokens, tokenizer) expected_text = ( "Sure, here's a short paragraph about emoji, " "where each word is followed by an emoji:\n\n" "Hey there! 👋 I'm here to help answer any questions you have about emoji" ) assert total_text == expected_text @require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") def test_stop_str_handler_not_stop( llama_tokenizer_path: str, ): stop_strs = ["^^"] tokenizer = Tokenizer(llama_tokenizer_path) stop_handler = StopStrHandler(stop_strs, tokenizer) total_text = stop_handler_process_tokens(stop_handler, para_input_tokens, tokenizer) assert total_text == DECODED_PARAGRAPH @require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") def test_stop_str_handler_return_cached_tokens( llama_tokenizer_path: str, ): tokens = para_input_tokens[:26] # until "\n\n" stop_strs = ["\n\n\n"] tokenizer = Tokenizer(llama_tokenizer_path) stop_handler = StopStrHandler(stop_strs, tokenizer) total_text = stop_handler_process_tokens(stop_handler, tokens, tokenizer) expected_text = ( "Sure, here's a short paragraph about emoji, where each word is followed by an emoji:\n\n" ) assert total_text == expected_text @require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") def test_stop_str_handler_throughput( llama_tokenizer_path: str, ): stop_strs = ["[INST]"] tokenizer = Tokenizer(llama_tokenizer_path) stop_handler = StopStrHandler(stop_strs, tokenizer) tokens = para_input_tokens * 20 returned_tokens = [] tbegin = time.perf_counter() for token in tokens: returned_tokens += stop_handler.put(token) assert not stop_handler.stop_triggered tend = time.perf_counter() throughput = len(tokens) / (tend - tbegin) print( f"num tokens = {len(tokens)}, " f"time elapsed = {tend - tbegin:.5f} sec, " f"throughput = {throughput}" ) assert throughput >= 100000 emoji_tokens_expected_result = [ # HF: "�����", SentencePiece: "�👀" ([177, 243, 162, 148, 131], ("�����", "�👀")), # Both: "👀👀" ([243, 162, 148, 131, 243, 162, 148, 131], ("👀👀",)), # Both: "👀👀👀" ([243, 162, 148, 131, 243, 162, 148, 131, 243, 162, 148, 131], ("👀👀👀",)), # HF: "👀�������", SentencePiece: "👀���👀" ([243, 162, 148, 131, 162, 148, 131, 243, 162, 148, 131], ("👀�������", "👀���👀")), # Both: "👀��� have👀" ([243, 162, 148, 131, 162, 148, 131, 505, 243, 162, 148, 131], ("👀��� have👀",)), ] @pytest.mark.parametrize("tokens_and_results", emoji_tokens_expected_result) @require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") def test_text_streamer_emojis( llama_tokenizer_path: str, tokens_and_results: Tuple[List[int], Tuple[str]], # noqa: UP006 ): text_streamer = TextStreamer(Tokenizer(llama_tokenizer_path)) total_text = "" tokens, expected_results = tokens_and_results for token in tokens: total_text += text_streamer.put([token]) total_text += text_streamer.finish() assert total_text in expected_results if __name__ == "__main__": test_text_streamer() test_stop_str_handler_stop() test_stop_str_handler_not_stop() test_stop_str_handler_return_cached_tokens() test_stop_str_handler_throughput() for tokens_and_res in emoji_tokens_expected_result: test_text_streamer_emojis(tokens_and_res)