Temporary fix for input truncation

IBM · Mar 28, 2024 · cde0474 · cde0474
1 parent 029c685
commit cde0474
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/vllm/entrypoints/grpc/grpc_server.py b/vllm/entrypoints/grpc/grpc_server.py
@@ -461,13 +461,26 @@ async def _validate_prompt_and_tokenize(
         prompt: Optional[str],
         context: ServicerContext,
     ) -> Tuple[List[int], bool]:
-        tokenize_kwargs = {"truncation": True,
-                           "max_length": truncate_input_tokens} \
-            if truncate_input_tokens is not None else {}
-
         max_model_len = self.config.max_model_len
+        # tokenize_kwargs = {"truncation": True,
+        #                    "max_length": truncate_input_tokens} \
+        #     if truncate_input_tokens is not None else {
+        #       "truncation": True, "max_length": max_model_len + 1}
+        tokenize_kwargs = {}
+
         input_ids = await self.tokenizer_group.encode_async(
             prompt, **tokenize_kwargs)
+
+        #TODO this is temporary until truncation option is added
+        # to the TokenizerGroup encode methods
+        if truncate_input_tokens and truncate_input_tokens < len(input_ids):
+            input_ids = input_ids[-truncate_input_tokens:]
+            if not sampling_params.skip_special_tokens:
+                add_bos_token = getattr(self.tokenizer, "add_bos_token", False)
+                if add_bos_token:
+                    input_ids[0] = self.tokenizer.bos_token_id
+        # -----------------------------------------------
+
         token_num = len(input_ids)
 
         if token_num >= max_model_len:

diff --git a/vllm/tgis_utils/args.py b/vllm/tgis_utils/args.py
@@ -92,7 +92,8 @@ def postprocess_tgis_args(args: argparse.Namespace) -> argparse.Namespace:
         args.dtype = args.dtype_str
     if args.quantize:
         if args.quantization and args.quantization != args.quantize:
-            raise ValueError("Inconsistent quantize and quantization arg values")
+            raise ValueError(
+                "Inconsistent quantize and quantization arg values")
         args.quantization = args.quantize
     if args.num_gpus is not None or args.num_shard is not None:
         if args.num_gpus is not None and args.num_shard is not None \