wehub-resource-sync
286 lines
9.5 KiB
Python
286 lines
9.5 KiB
Python
import asyncio
|
|
from typing import List # noqa: UP035
|
|
|
|
from mlc_llm.protocol.generation_config import GenerationConfig
|
|
from mlc_llm.serve import AsyncMLCEngine, EngineConfig
|
|
from mlc_llm.testing import require_test_model
|
|
|
|
prompts = [
|
|
"What is the meaning of life?",
|
|
"Introduce the history of Pittsburgh to me. Please elaborate in detail.",
|
|
"Write a three-day Seattle travel plan. Please elaborate in detail.",
|
|
"What is Alaska famous of? Please elaborate in detail.",
|
|
"What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501
|
|
"What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
|
|
"Why is Vitamin D important to human beings? Please elaborate in detail.",
|
|
"Where is milk tea originated from? Please elaborate in detail.",
|
|
"Where is the southernmost place in United States? Please elaborate in detail.",
|
|
"Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501
|
|
]
|
|
|
|
|
|
@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
|
|
async def test_engine_generate(model: str):
|
|
# Create engine
|
|
async_engine = AsyncMLCEngine(
|
|
model=model,
|
|
mode="server",
|
|
engine_config=EngineConfig(max_total_sequence_length=4096),
|
|
)
|
|
|
|
num_requests = 10
|
|
max_tokens = 256
|
|
generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7)
|
|
|
|
output_texts: List[List[str]] = [ # noqa: UP006
|
|
["" for _ in range(generation_cfg.n)] for _ in range(num_requests)
|
|
]
|
|
|
|
async def generate_task(
|
|
async_engine: AsyncMLCEngine,
|
|
prompt: str,
|
|
generation_cfg: GenerationConfig,
|
|
request_id: str,
|
|
):
|
|
print(f"generate task for request {request_id}")
|
|
rid = int(request_id)
|
|
async for delta_outputs in async_engine._generate(
|
|
prompt, generation_cfg, request_id=request_id
|
|
):
|
|
if len(delta_outputs) == generation_cfg.n:
|
|
for i, delta_output in enumerate(delta_outputs):
|
|
output_texts[rid][i] += delta_output.delta_text
|
|
else:
|
|
assert len(delta_outputs) == 1
|
|
assert len(delta_outputs[0].request_final_usage_json_str) != 0
|
|
|
|
tasks = [
|
|
asyncio.create_task(
|
|
generate_task(async_engine, prompts[i], generation_cfg, request_id=str(i))
|
|
)
|
|
for i in range(num_requests)
|
|
]
|
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
# Print output.
|
|
print("All finished")
|
|
for req_id, outputs in enumerate(output_texts):
|
|
print(f"Prompt {req_id}: {prompts[req_id]}")
|
|
if len(outputs) == 1:
|
|
print(f"Output {req_id}:{outputs[0]}\n")
|
|
else:
|
|
for i, output in enumerate(outputs):
|
|
print(f"Output {req_id}({i}):{output}\n")
|
|
|
|
async_engine.terminate()
|
|
del async_engine
|
|
|
|
|
|
@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
|
|
async def test_chat_completion(model: str):
|
|
# Create engine
|
|
async_engine = AsyncMLCEngine(
|
|
model=model,
|
|
mode="server",
|
|
engine_config=EngineConfig(max_total_sequence_length=4096),
|
|
)
|
|
|
|
num_requests = 2
|
|
max_tokens = 32
|
|
n = 1
|
|
output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006
|
|
|
|
async def generate_task(prompt: str, request_id: str):
|
|
print(f"generate chat completion task for request {request_id}")
|
|
rid = int(request_id)
|
|
async for response in await async_engine.chat.completions.create( # noqa: F821
|
|
messages=[{"role": "user", "content": prompt}],
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
n=n,
|
|
request_id=request_id,
|
|
stream=True,
|
|
):
|
|
for choice in response.choices:
|
|
assert choice.delta.role == "assistant"
|
|
assert isinstance(choice.delta.content, str)
|
|
output_texts[rid][choice.index] += choice.delta.content
|
|
|
|
tasks = [
|
|
asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
|
|
for i in range(num_requests)
|
|
]
|
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
# Print output.
|
|
print("Chat completion all finished")
|
|
for req_id, outputs in enumerate(output_texts):
|
|
print(f"Prompt {req_id}: {prompts[req_id]}")
|
|
if len(outputs) == 1:
|
|
print(f"Output {req_id}:{outputs[0]}\n")
|
|
else:
|
|
for i, output in enumerate(outputs):
|
|
print(f"Output {req_id}({i}):{output}\n")
|
|
|
|
async_engine.terminate()
|
|
del async_engine
|
|
|
|
|
|
@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
|
|
async def test_chat_completion_non_stream(model: str):
|
|
# Create engine
|
|
async_engine = AsyncMLCEngine(
|
|
model=model,
|
|
mode="server",
|
|
engine_config=EngineConfig(max_total_sequence_length=4096),
|
|
)
|
|
|
|
num_requests = 2
|
|
max_tokens = 32
|
|
n = 1
|
|
output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006
|
|
|
|
async def generate_task(prompt: str, request_id: str):
|
|
print(f"generate chat completion task for request {request_id}")
|
|
rid = int(request_id)
|
|
response = await async_engine.chat.completions.create( # noqa: F821
|
|
messages=[{"role": "user", "content": prompt}],
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
n=n,
|
|
request_id=request_id,
|
|
)
|
|
for choice in response.choices:
|
|
assert choice.message.role == "assistant"
|
|
assert isinstance(choice.message.content, str)
|
|
output_texts[rid][choice.index] += choice.message.content
|
|
|
|
tasks = [
|
|
asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
|
|
for i in range(num_requests)
|
|
]
|
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
# Print output.
|
|
print("Chat completion all finished")
|
|
for req_id, outputs in enumerate(output_texts):
|
|
print(f"Prompt {req_id}: {prompts[req_id]}")
|
|
if len(outputs) == 1:
|
|
print(f"Output {req_id}:{outputs[0]}\n")
|
|
else:
|
|
for i, output in enumerate(outputs):
|
|
print(f"Output {req_id}({i}):{output}\n")
|
|
|
|
async_engine.terminate()
|
|
del async_engine
|
|
|
|
|
|
@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
|
|
async def test_completion(model: str):
|
|
# Create engine
|
|
async_engine = AsyncMLCEngine(
|
|
model=model,
|
|
mode="server",
|
|
engine_config=EngineConfig(max_total_sequence_length=4096),
|
|
)
|
|
|
|
num_requests = 2
|
|
max_tokens = 128
|
|
n = 1
|
|
output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006
|
|
|
|
async def generate_task(prompt: str, request_id: str):
|
|
print(f"generate completion task for request {request_id}")
|
|
rid = int(request_id)
|
|
async for response in await async_engine.completions.create( # noqa: F821
|
|
prompt=prompt,
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
n=n,
|
|
request_id=request_id,
|
|
stream=True,
|
|
extra_body={"debug_config": {"ignore_eos": True}},
|
|
):
|
|
for choice in response.choices:
|
|
output_texts[rid][choice.index] += choice.text
|
|
|
|
tasks = [
|
|
asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
|
|
for i in range(num_requests)
|
|
]
|
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
# Print output.
|
|
print("Completion all finished")
|
|
for req_id, outputs in enumerate(output_texts):
|
|
print(f"Prompt {req_id}: {prompts[req_id]}")
|
|
if len(outputs) == 1:
|
|
print(f"Output {req_id}:{outputs[0]}\n")
|
|
else:
|
|
for i, output in enumerate(outputs):
|
|
print(f"Output {req_id}({i}):{output}\n")
|
|
|
|
async_engine.terminate()
|
|
del async_engine
|
|
|
|
|
|
@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
|
|
async def test_completion_non_stream(model: str):
|
|
# Create engine
|
|
async_engine = AsyncMLCEngine(
|
|
model=model,
|
|
mode="server",
|
|
engine_config=EngineConfig(max_total_sequence_length=4096),
|
|
)
|
|
|
|
num_requests = 2
|
|
max_tokens = 128
|
|
n = 1
|
|
output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006
|
|
|
|
async def generate_task(prompt: str, request_id: str):
|
|
print(f"generate completion task for request {request_id}")
|
|
rid = int(request_id)
|
|
response = await async_engine.completions.create( # noqa: F821
|
|
prompt=prompt,
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
n=n,
|
|
request_id=request_id,
|
|
extra_body={"debug_config": {"ignore_eos": True}},
|
|
)
|
|
for choice in response.choices:
|
|
output_texts[rid][choice.index] += choice.text
|
|
|
|
tasks = [
|
|
asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
|
|
for i in range(num_requests)
|
|
]
|
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
# Print output.
|
|
print("Completion all finished")
|
|
for req_id, outputs in enumerate(output_texts):
|
|
print(f"Prompt {req_id}: {prompts[req_id]}")
|
|
if len(outputs) == 1:
|
|
print(f"Output {req_id}:{outputs[0]}\n")
|
|
else:
|
|
for i, output in enumerate(outputs):
|
|
print(f"Output {req_id}({i}):{output}\n")
|
|
|
|
async_engine.terminate()
|
|
del async_engine
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_engine_generate())
|
|
asyncio.run(test_chat_completion())
|
|
asyncio.run(test_chat_completion_non_stream())
|
|
asyncio.run(test_completion())
|
|
asyncio.run(test_completion_non_stream())
|