Add option to completion API to truncate prompt tokens (#3144)

2024-04-05 19:15:42 +02:00
parent cfaf49a167
commit 1d7c940d74
4 changed files with 41 additions and 8 deletions
--- a/vllm/entrypoints/openai/serving_completion.py
+++ b/vllm/entrypoints/openai/serving_completion.py
@@ -137,10 +137,16 @@ class OpenAIServingCompletion(OpenAIServing):
            for i, prompt in enumerate(prompts):
                if prompt_is_tokens:
                    input_ids = self._validate_prompt_and_tokenize(
-                        request, prompt_ids=prompt)
+                        request,
+                        prompt_ids=prompt,
+                        truncate_prompt_tokens=sampling_params.
+                        truncate_prompt_tokens)
                else:
                    input_ids = self._validate_prompt_and_tokenize(
-                        request, prompt=prompt)
+                        request,
+                        prompt=prompt,
+                        truncate_prompt_tokens=sampling_params.
+                        truncate_prompt_tokens)

                generators.append(
                    self.engine.generate(prompt,