wehub-resource-sync
40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
"""Mock testing engine I/O conventions
|
|
|
|
Mock test only can help checking the overall input
|
|
output processing options are passed correctly
|
|
"""
|
|
|
|
import pytest
|
|
import tvm
|
|
|
|
from mlc_llm.serve import MLCEngine
|
|
from mlc_llm.testing import require_test_model
|
|
|
|
# test category "unittest"
|
|
pytestmark = [pytest.mark.unittest]
|
|
|
|
|
|
# NOTE: we only need tokenizers in folder
|
|
# launch time of mock test is fast so we can put it in unittest
|
|
@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC")
|
|
def test_completion_api(model: str):
|
|
engine = MLCEngine(model, tvm.cpu(), model_lib="mock://echo")
|
|
param_dict = {
|
|
"top_p": 0.6,
|
|
"temperature": 0.9,
|
|
"frequency_penalty": 0.1,
|
|
"presence_penalty": 0.1,
|
|
"n": 2,
|
|
}
|
|
response = engine.chat.completions.create(
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
**param_dict,
|
|
)
|
|
# echo mock will echo back the generation config
|
|
for k, v in param_dict.items():
|
|
assert response.usage.extra[k] == v
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_completion_api()
|