[python] add max_logprobs vllm configuration to EngineArgs (#2154)

deepjavalibrary · Jul 10, 2024 · 823563f · 823563f
1 parent 486cf71
commit 823563f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/engines/python/setup/djl_python/properties_manager/lmi_dist_rb_properties.py b/engines/python/setup/djl_python/properties_manager/lmi_dist_rb_properties.py
@@ -51,6 +51,7 @@ class LmiDistRbProperties(Properties):
     max_lora_rank: Optional[int] = 16
     lora_extra_vocab_size: Optional[int] = 256
     max_cpu_loras: Optional[int] = None
+    max_logprobs: Optional[int] = 5
 
     @model_validator(mode='after')
     def validate_mpi(self):

diff --git a/engines/python/setup/djl_python/properties_manager/vllm_rb_properties.py b/engines/python/setup/djl_python/properties_manager/vllm_rb_properties.py
@@ -47,6 +47,7 @@ class VllmRbProperties(Properties):
     # Neuron vLLM properties
     device: Optional[str] = None
     preloaded_model: Optional[Any] = None
+    max_logprobs: Optional[int] = 5
 
     @field_validator('engine')
     def validate_engine(cls, engine):

diff --git a/engines/python/setup/djl_python/rolling_batch/rolling_batch_vllm_utils.py b/engines/python/setup/djl_python/rolling_batch/rolling_batch_vllm_utils.py
@@ -224,4 +224,5 @@ def get_engine_args_from_config(config: VllmRbProperties) -> EngineArgs:
             max_lora_rank=config.max_lora_rank,
             lora_extra_vocab_size=config.lora_extra_vocab_size,
             max_cpu_loras=config.max_cpu_loras,
-            revision=config.revision)
+            revision=config.revision,
+            max_logprobs=config.max_logprobs)