wehub-resource-sync
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
from mlc_llm.protocol.generation_config import GenerationConfig
|
|
from mlc_llm.serve import data
|
|
from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
|
|
|
|
|
|
def get_test_image(config) -> data.ImageData:
|
|
return data.ImageData.from_url("https://llava-vl.github.io/static/images/view.jpg", config)
|
|
|
|
|
|
def test_engine_generate():
|
|
# Create engine
|
|
model = "dist/llava-1.5-7b-hf-q4f16_1-MLC/params"
|
|
model_lib = "dist/llava-1.5-7b-hf-q4f16_1-MLC/llava-1.5-7b-hf-q4f16_1-MLC.so"
|
|
engine = SyncMLCEngine(
|
|
model=model,
|
|
model_lib=model_lib,
|
|
mode="server",
|
|
engine_config=EngineConfig(max_total_sequence_length=4096),
|
|
)
|
|
max_tokens = 256
|
|
|
|
with open(Path(model) / "mlc-chat-config.json", encoding="utf-8") as file:
|
|
model_config = json.load(file)
|
|
|
|
prompts = [
|
|
[
|
|
data.TextData("USER: "),
|
|
get_test_image(model_config),
|
|
data.TextData("\nWhat does this image represent? ASSISTANT:"),
|
|
],
|
|
[
|
|
data.TextData("USER: "),
|
|
get_test_image(model_config),
|
|
data.TextData("\nIs there a dog in this image? ASSISTANT:"),
|
|
],
|
|
[data.TextData("USER: What is the meaning of life? ASSISTANT:")],
|
|
]
|
|
|
|
output_texts, _ = engine.generate(
|
|
prompts, GenerationConfig(max_tokens=max_tokens, stop_token_ids=[2])
|
|
)
|
|
|
|
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")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_engine_generate()
|