wehub-resource-sync
85 lines
2.9 KiB
Python
85 lines
2.9 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-q0f16-MLC",
|
|
"Llama-2-7b-chat-hf-q4f16_1-MLC",
|
|
)
|
|
async def test_engine_generate(model: str, small_model: str):
|
|
# Create engine
|
|
async_engine = AsyncMLCEngine(
|
|
model=model,
|
|
mode="server",
|
|
engine_config=EngineConfig(
|
|
additional_models=[small_model],
|
|
speculative_mode="small_draft",
|
|
),
|
|
)
|
|
|
|
num_requests = 10
|
|
max_tokens = 256
|
|
generation_cfg = GenerationConfig(max_tokens=max_tokens)
|
|
|
|
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
|
|
):
|
|
assert len(delta_outputs) == generation_cfg.n
|
|
for i, delta_output in enumerate(delta_outputs):
|
|
output_texts[rid][i] += delta_output.delta_text
|
|
|
|
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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_engine_generate())
|