logprobs for vllm >= 0.5.5

turbo-llm · Nov 12, 2024 · b4c717a · b4c717a
1 parent a66986f
commit b4c717a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/turbo_alignment/generators/vllm_chat.py b/turbo_alignment/generators/vllm_chat.py
@@ -51,6 +51,7 @@ def __init__(
             skip_special_tokens=custom_generation_settings.skip_special_tokens,
             stop_token_ids=eos_token_id,
             max_tokens=transformers_settings.max_new_tokens,
+            logprobs=transformers_settings.logprobs,
             **beam_search_params,
         )
         self._lora_request = lora_request

diff --git a/turbo_alignment/settings/tf/generation.py b/turbo_alignment/settings/tf/generation.py
@@ -11,3 +11,4 @@ class GeneratorTransformersSettings(ExtraFieldsNotAllowedBaseModel):
     top_k: int = 50
     temperature: float = 1.0
     stop_strings: str | list[str] = '</s>'
+    logprobs: int = 1