wehub-resource-sync
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
import base64
|
|
from typing import Dict, List, Optional # noqa: UP035
|
|
|
|
import requests
|
|
|
|
from mlc_llm.json_ffi import JSONFFIEngine
|
|
from mlc_llm.testing import require_test_model
|
|
|
|
|
|
def base64_encode_image(url: str) -> str:
|
|
response = requests.get(url)
|
|
response.raise_for_status() # Ensure we got a successful response
|
|
image_data = base64.b64encode(response.content)
|
|
image_data_str = image_data.decode("utf-8")
|
|
data_url = f"data:image/jpeg;base64,{image_data_str}"
|
|
return data_url
|
|
|
|
|
|
image_prompts = [
|
|
[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": f"{base64_encode_image('https://llava-vl.github.io/static/images/view.jpg')}",
|
|
},
|
|
{"type": "text", "text": "What does the image represent?"},
|
|
],
|
|
}
|
|
]
|
|
]
|
|
|
|
|
|
def run_chat_completion(
|
|
engine: JSONFFIEngine,
|
|
model: str,
|
|
prompts: List[List[Dict]] = image_prompts, # noqa: UP006
|
|
tools: Optional[List[Dict]] = None, # noqa: UP006
|
|
):
|
|
num_requests = 1
|
|
max_tokens = 64
|
|
n = 1
|
|
output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006
|
|
|
|
for rid in range(num_requests):
|
|
print(f"chat completion for request {rid}")
|
|
for response in engine.chat.completions.create(
|
|
messages=prompts[rid],
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
n=n,
|
|
request_id=str(rid),
|
|
tools=tools,
|
|
):
|
|
for choice in response.choices:
|
|
assert choice.delta.role == "assistant"
|
|
assert isinstance(choice.delta.content[0], Dict) # noqa: UP006
|
|
assert choice.delta.content[0]["type"] == "text"
|
|
output_texts[rid][choice.index] += choice.delta.content[0]["text"]
|
|
|
|
# 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")
|
|
|
|
|
|
@require_test_model("llava-1.5-7b-hf-q4f16_1-MLC")
|
|
def test_chat_completion():
|
|
# Create engine.
|
|
engine = JSONFFIEngine(
|
|
model, # noqa: F821
|
|
max_total_sequence_length=1024,
|
|
)
|
|
|
|
run_chat_completion(engine, model) # noqa: F821
|
|
|
|
# Test malformed requests.
|
|
for response in engine._raw_chat_completion("malformed_string", n=1, request_id="123"):
|
|
assert len(response.choices) == 1
|
|
assert response.choices[0].finish_reason == "error"
|
|
|
|
engine.terminate()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_chat_completion()
|