[Misc]Add customized information for models (#4132)

This commit is contained in:
Jee Li
2024-05-01 12:18:14 +08:00
committed by GitHub
parent c3845d82dc
commit d6f4bd7cdd
9 changed files with 87 additions and 0 deletions

View File

@ -49,3 +49,18 @@ def test_models(
f"Test{i}:\nHF: {hf_output_str!r}\nvLLM: {vllm_output_str!r}")
assert hf_output_ids == vllm_output_ids, (
f"Test{i}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}")
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["float"])
def test_model_print(
vllm_runner,
model: str,
dtype: str,
) -> None:
vllm_model = vllm_runner(model, dtype=dtype)
# This test is for verifying whether the model's extra_repr
# can be printed correctly.
print(vllm_model.model.llm_engine.model_executor.driver_worker.
model_runner.model)
del vllm_model