[ChatStream] Transformer Mock: Mocking Transformer Responses
Hello from the Product Development Team at Qualiteg.
In this article we explain how to create mock data. This is for a feature formally called "Transformer Mock", which records actual LLM output and replays it later.
Why would you need such a thing? It is used for testing LLM applications (unit tests and the like). In classic unit testing, the premise is that the expected output for a given input is fixed.
LLMs, however, by their very nature return a different response every time even for the same input. That is one of the appealing things about generative AI, but it becomes a headache when you want to write classic unit tests.
At this point, astute readers may be thinking, "If you want the same output for the same input, just fix the seed." It turns out, though, that even when you fix the seed, fix the input, and fix all of the sampling parameters, you can still get different output on a different type of GPU.
That would mean unit tests start failing the moment you switch GPUs, which is a real problem. So the idea behind this feature is: record the values fed into and produced by a given GPU, and at unit-test time "replay" that recording, thereby simulating the GPU's computation inputs and results.
This lets unit tests exercise a large portion of the ChatStream code (higher coverage per test run), which improves the reliability of the unit tests.
In addition, loading a large model can take tens of minutes, so even with a nightly CI build, a lot of time gets eaten up by parts that are not essential (and do not need to be covered). Emulation with this feature can dramatically shorten that time as well.
How to Create Mock Data
This section describes Mock mode, which lets you produce the same responses a model would give without actually loading the model.
What is Transformer Mock mode?
By recording pairs of inputs to and outputs from the Model and Tokenizer in advance and then replaying them,
you can make the system behave as if a Model and Tokenizer were present even when they are not.
Emulating the Model and Tokenizer in this way is what we call Mock mode.
Benefits of Transformer Mock mode
- No time spent loading model data.
- Reproducible output (AI assistant responses).
This makes it easy to set up evaluations and tests of everything other than the model itself.
Difference from Generator Mock
A similar feature is Generator Mock.
Whereas Transformer Mock mode records and replays the behavior of the actual Model and Tokenizer, Generator Mock
responds with dummy text after receiving input. Transformer Mock mode only accepts the predetermined inputs, while Generator Mock responds with dummy text to any input whatsoever.
Generator Mock is useful for things like checking API behavior, but the coverage achieved when running test code is considerably lower than with Transformer Mock mode. If coverage matters to you, Transformer Mock mode is recommended.
Recording and Replay
Recording for Transformer Mock mode: Probe mode
Strictly speaking, replaying the behavior of the Model and Tokenizer is called Transformer Mock mode,
and the mode that records the behavior of the Model and Tokenizer is called Probe mode.
Setting probe_mode_enabled=True as shown below enables Probe mode.
chat_stream = ChatStream(
num_of_concurrent_executions=2,
max_queue_size=5,
model=model,
tokenizer=tokenizer,
num_gpus=num_gpus,
device=device,
chat_prompt_clazz=ChatPrompt,
add_special_tokens=False,
max_new_tokens=128,
context_len=1024,
temperature=0.7,
top_k=10,
client_roles=client_role_free_access,
locale='ja',
token_sampler=TokenSamplerIsok(),
seed=42,
probe_mode_enabled=True,
)
Start the ChatStream server with probe_mode_enabled=True, enter text from the UI, and generate responses.
Simply chatting as usual like this automatically records the inputs and responses.
The recorded data is saved in the following directory:
[home_dir]/.cache/chatstream/probe_data
Emulating the Model and Tokenizer in Transformer Mock mode
Using MockTransformer, you can emulate the Model and Tokenizer with the recorded data.
MockTransformer(parent_dir_path=[parent directory], dirname=[name of the directory where the recorded data was saved],
wait_sec=[wait applied each time one token is generated (seconds)])
If [parent directory] is omitted,
[home_dir]/.cache/chatstream/probe_data
is used as the directory.
Sample code
mock_transformer = MockTransformer(parent_dir_path=mock_data_dir, dirname=mock_data_name, wait_sec=0)
model = mock_transformer.get_model() # model
tokenizer = mock_transformer.get_tokenizer() # tokenizer
token_sampler = mock_transformer.get_token_sampler() # sampling class
if device.type == 'cuda' and num_gpus == 1:
model.to(device)
chat_stream = ChatStream(
num_of_concurrent_executions=2,
max_queue_size=5,
model=model,
tokenizer=tokenizer,
num_gpus=num_gpus,
device=device,
chat_prompt_clazz=ChatPrompt,
add_special_tokens=False,
max_new_tokens=128, # The maximum size of the newly generated tokens
context_len=1024, # The size of the context (in terms of the number of tokens)
temperature=0.7, # The temperature value for randomness in prediction
top_k=10, # Value of top K for sampling
top_p=0.9, # Value of top P for sampling,
# repetition_penalty=1.05,
client_roles=client_role_free_access,
locale='ja',
token_sampler=token_sampler,
)
Starting the ChatStream server with this configuration runs it in Transformer Mock mode.
Note
The text you can enter, and its order, must be the same text and order as when the recording was made.