change input

Author: LunrEclipse

vllm/engine/async_llm_engine.py

@@ -1058,7 +1058,7 @@ async def beam_search(
         tokenizedLength = len(tokenizedPrompt)
 
         sort_beams_key = create_sort_beams_key_function(
-            tokenizer, length_penalty=length_penalty)
+            tokenizer.eos_token_id, length_penalty)
 
         beam_search_params = SamplingParams(logprobs=2 * beam_width,
                                             max_tokens=1,

vllm/engine/multiprocessing/client.py

@@ -463,7 +463,7 @@ async def beam_search(
         tokenizedLength = len(tokenizedPrompt)
 
         sort_beams_key = create_sort_beams_key_function(
-            tokenizer, length_penalty=length_penalty)
+            tokenizer.eos_token_id, length_penalty)
 
         beam_search_params = SamplingParams(logprobs=2 * beam_width,
                                             max_tokens=1,

vllm/utils.py

@@ -1414,10 +1414,10 @@ def get_beam_search_score(
     return cumulative_logprob / (seq_len**length_penalty)
 
 
-def create_sort_beams_key_function(tokenizer, length_penalty):
+def create_sort_beams_key_function(eos_token_id: int, length_penalty):
 
     def sort_beams_key(x: BeamSearchSequence) -> float:
-        return get_beam_search_score(x.tokens, x.cum_logprob,
-                                     tokenizer.eos_token_id, length_penalty)
+        return get_beam_search_score(x.tokens, x.cum_logprob, eos_token_id,
+                                     length_penalty)
 
     return sort_beams_key