update tests

vllm-project · Aug 26, 2024 · 1661f7c · 1661f7c
1 parent 599afa2
commit 1661f7c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/tests/e2e/vLLM/__init__.py b/tests/e2e/vLLM/__init__.py
diff --git 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
-