Tool completion respect num new tokens (#469)

* Tool completion respect num new tokens

* more fix

* remove unused import

* format

* empty

* no cover

Author: yunfeng-scale

model-engine/model_engine_server/inference/batch_inference/vllm_batch.py

@@ -185,7 +185,7 @@ def __repr__(self) -> str:
 
         outputs = await generate_with_vllm(
             llm,
-            content.max_new_tokens,
+            [generations[iter[1]].remaining_tokens for iter in iter_prompts],
             content.temperature,
             content.stop_sequences,
             content.return_token_log_probs,
@@ -260,7 +260,10 @@ def tool_func(text: str, past_context: Optional[str]):
             gen_item.generated_text += new_text
 
             # If we didn't just execute a tool, we're done
-            if not gen_item.generated_text.endswith(tool.tool_context_end):
+            if (
+                not gen_item.generated_text.endswith(tool.tool_context_end)
+                or gen_item.remaining_tokens <= 0
+            ):
                 gen_item.completed = True
                 continue
 
@@ -316,7 +319,7 @@ async def batch_inference():
 
         outputs = await generate_with_vllm(
             llm,
-            content.max_new_tokens,
+            [content.max_new_tokens] * len(prompts),
             content.temperature,
             content.stop_sequences,
             content.return_token_log_probs,
@@ -358,24 +361,23 @@ async def generate_with_vllm(
     top_p,
     prompts,
     bar,
-) -> List[CompletionOutput]:
+) -> List[CompletionOutput]:  # pragma: no cover
     from vllm import SamplingParams
 
     # Add the requests to the engine.
-    sampling_params = SamplingParams(
-        max_tokens=max_new_tokens,
-        temperature=temperature,
-        stop=stop_sequences,
-        logprobs=1 if return_token_log_probs else None,
-        presence_penalty=presence_penalty or 0.0,
-        frequency_penalty=frequency_penalty or 0.0,
-        top_k=top_k or -1,
-        top_p=top_p or 1.0,
-    )
-
     results_generators = []
-    for prompt in prompts:
+    for idx, prompt in enumerate(prompts):
         request_id = random_uuid()
+        sampling_params = SamplingParams(
+            max_tokens=max_new_tokens[idx],
+            temperature=temperature,
+            stop=stop_sequences,
+            logprobs=1 if return_token_log_probs else None,
+            presence_penalty=presence_penalty or 0.0,
+            frequency_penalty=frequency_penalty or 0.0,
+            top_k=top_k or -1,
+            top_p=top_p or 1.0,
+        )
         results_generator = await engine.add_request(
             request_id, prompt, sampling_params, None, time.monotonic()
         )