File size: 2,156 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | from __future__ import annotations
from diffulex.sampling_params import SamplingParams
from diffulex.server.protocol import (
ChatInput,
PromptInput,
ServingDelta,
ServingGenerate,
ServingReply,
serving_command_from_dict,
serving_command_to_dict,
serving_event_from_dict,
serving_event_to_dict,
)
def test_serving_generate_protocol_round_trips_chat_input():
command = ServingGenerate(
rid="rid-1",
input=ChatInput([{"role": "user", "content": "hi"}]),
sampling_params=SamplingParams(max_tokens=7, temperature=0.0, max_nfe=16),
stream=True,
stream_mode="block_append",
user="debug-user",
created_time=123.0,
)
decoded = serving_command_from_dict(serving_command_to_dict(command))
assert isinstance(decoded, ServingGenerate)
assert decoded.rid == "rid-1"
assert isinstance(decoded.input, ChatInput)
assert decoded.input.messages == [{"role": "user", "content": "hi"}]
assert decoded.sampling_params.max_tokens == 7
assert decoded.sampling_params.max_nfe == 16
assert decoded.stream is True
assert decoded.user == "debug-user"
def test_serving_generate_protocol_round_trips_prompt_ids():
command = ServingGenerate(
rid="rid-2",
input=PromptInput([1, 2, 3]),
sampling_params=SamplingParams(max_tokens=3),
)
decoded = serving_command_from_dict(serving_command_to_dict(command))
assert isinstance(decoded, ServingGenerate)
assert isinstance(decoded.input, PromptInput)
assert decoded.input.prompt == [1, 2, 3]
assert decoded.stream_mode == "denoise"
assert decoded.sampling_params.max_nfe == 512
def test_serving_event_protocol_round_trips_reply_and_delta():
reply = ServingReply(rid="rid-1", text="ok", token_ids=[1, 2], nfe=3, finish_reason="stop")
delta = ServingDelta(rid="rid-1", token_offset=0, text="o", token_ids=[1], nfe=1)
decoded_reply = serving_event_from_dict(serving_event_to_dict(reply))
decoded_delta = serving_event_from_dict(serving_event_to_dict(delta))
assert decoded_reply == reply
assert decoded_delta == delta
|