wehub-resource-sync
61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
from typing import List # noqa: UP035
|
|
|
|
from mlc_llm.protocol.generation_config import GenerationConfig
|
|
from mlc_llm.serve import EngineConfig, MLCEngine
|
|
|
|
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
|
|
]
|
|
|
|
|
|
def test_engine_generate() -> None:
|
|
engine = MLCEngine(
|
|
model="dist/rwkv-6-world-1b6-q0f16-MLC",
|
|
model_lib="dist/rwkv-6-world-1b6-q0f16-MLC/rwkv-6-world-1b6-q0f16-MLC-cuda.so",
|
|
mode="server",
|
|
engine_config=EngineConfig(
|
|
max_num_sequence=8,
|
|
max_history_size=1,
|
|
),
|
|
)
|
|
|
|
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)
|
|
]
|
|
for rid in range(num_requests):
|
|
print(f"generating for request {rid}")
|
|
for delta_outputs in engine._generate(prompts[rid], generation_cfg, request_id=str(rid)):
|
|
assert len(delta_outputs) == generation_cfg.n
|
|
for i, delta_output in enumerate(delta_outputs):
|
|
output_texts[rid][i] += delta_output.delta_text
|
|
|
|
# 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")
|
|
|
|
engine.terminate()
|
|
del engine
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_engine_generate()
|