From 1661f7c9701edde4c5dcd809d7fbf7c54089da75 Mon Sep 17 00:00:00 2001 From: Dipika Date: Mon, 26 Aug 2024 22:09:28 +0000 Subject: [PATCH] update tests --- tests/e2e/vLLM/__init__.py | 0 tests/e2e/vLLM/test_vllm.py | 11 +++++------ 2 files changed, 5 insertions(+), 6 deletions(-) create mode 100644 tests/e2e/vLLM/__init__.py diff --git a/tests/e2e/vLLM/__init__.py b/tests/e2e/vLLM/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/e2e/vLLM/test_vllm.py b/tests/e2e/vLLM/test_vllm.py index ec7338e62..71719287f 100644 --- a/tests/e2e/vLLM/test_vllm.py +++ b/tests/e2e/vLLM/test_vllm.py @@ -1,13 +1,13 @@ +import unittest + from transformers import AutoTokenizer +from vllm import LLM, SamplingParams from llmcompressor.modifiers.quantization import QuantizationModifier from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot -from vllm import LLM, SamplingParams - -import unittest - from tests.testing_utils import requires_gpu + @requires_gpu class TestvLLM(unittest.TestCase): MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct" @@ -43,8 +43,7 @@ class TestvLLM(unittest.TestCase): tokenizer.save_pretrained(SAVE_DIR) sampling_params = SamplingParams(temperature=0.80, top_p=0.95) - llm = LLM(model=model_path) + llm = LLM(model=SAVE_DIR) outputs = llm.generate(prompts, sampling_params) print(outputs) assert output -