fix dependencies.

vllm-project · Sep 23, 2024 · 5b15061 · 5b15061
1 parent cb14592
commit 5b15061
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/vllm/model_executor/model_loader/neuron.py b/vllm/model_executor/model_loader/neuron.py
@@ -14,8 +14,8 @@
 from vllm.model_executor.layers.quantization import get_quantization_config
 from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
 from vllm.model_executor.sampling_metadata import SamplingMetadata
-from vllm.sequence import SamplerOutput, SequenceOutput, CompletionSequenceGroupOutput, Logprob
-from transformers_neuronx.config import GenerationConfig
+from vllm.sequence import SequenceOutput, CompletionSequenceGroupOutput, Logprob
+from vllm.model_executor.layers.sampler import SamplerOutput
 
 TORCH_DTYPE_TO_NEURON_AMP = {
     "auto": "f32",

diff --git a/vllm/worker/neuron_model_runner.py b/vllm/worker/neuron_model_runner.py
@@ -281,7 +281,7 @@ def execute_model(
                 "NeuronModelRunner does not support multi-step execution.")
 
         # Update Neuron's generation configs from sampling_metadata
-        self._update_neuron_generation_config(ModelInputForNeuron.sampling_metadata)
+        self._update_neuron_generation_config(model_input.sampling_metadata)
 
         hidden_states = self.model(
             input_ids=model_input.input_tokens,