Add type annotations and make None the default for pooler arguments

Signed-off-by: Went-Liang <wenteng_liang@163.com>
vllm-project · Oct 29, 2024 · e62f65c · e62f65c
1 parent 2a1a748
commit e62f65c
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 12 deletions.
diff --git a/vllm/config.py b/vllm/config.py
@@ -162,7 +162,7 @@ def __init__(
             pooling_type: Optional[str] = None,
             pooling_norm: bool = False,
             pooling_softmax: bool = False,
-            pooling_step_tag_id: int = -1,
+            pooling_step_tag_id: Optional[int] = None,
             pooling_returned_token_ids: Optional[List[int]] = None) -> None:
         self.model = model
         self.tokenizer = tokenizer
@@ -270,9 +270,13 @@ def _init_multimodal_config(
         return None
 
     def _init_pooler_config(
-            self, pooling_type, pooling_norm, pooling_softmax,
-            pooling_step_tag_id,
-            pooling_returned_token_ids) -> Optional["PoolerConfig"]:
+        self,
+        pooling_type: Optional[str] = None,
+        pooling_norm: bool = False,
+        pooling_softmax: bool = False,
+        pooling_step_tag_id: Optional[int] = None,
+        pooling_returned_token_ids: Optional[List[int]] = None
+    ) -> Optional["PoolerConfig"]:
         if self.task == "embedding":
             return PoolerConfig(
                 pooling_type=pooling_type,
@@ -1694,7 +1698,7 @@ class PoolerConfig:
     pooling_type: Optional[str] = None
     pooling_norm: bool = False
     pooling_softmax: bool = False
-    pooling_step_tag_id: int = -1
+    pooling_step_tag_id: Optional[int] = None
     pooling_returned_token_ids: Optional[List[int]] = None
 
 

diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
@@ -188,7 +188,7 @@ class EngineArgs:
     pooling_type: Optional[str] = None
     pooling_norm: bool = False
     pooling_softmax: bool = False
-    pooling_step_tag_id: int = -1
+    pooling_step_tag_id: Optional[int] = None
     pooling_returned_token_ids: Optional[List[int]] = None
 
     def __post_init__(self):
@@ -860,7 +860,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument(
             '--pooling-type',
             choices=['LAST', 'ALL', 'CLS', 'STEP'],
-            default="LAST",
+            default=None,
             help='Used to configure the pooling method in the embedding model.'
         )
 
@@ -877,7 +877,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         parser.add_argument(
             '--pooling-step-tag-id',
             type=int,
-            default=-1,
+            default=None,
             help="When pooling-step-tag-id is not -1, it indicates "
             "that the score corresponding to the step-tag-ids in the "
             "generated sentence should be returned. Otherwise, it "

diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
@@ -162,7 +162,7 @@ def __init__(
         pooling_type: Optional[str] = None,
         pooling_norm: bool = False,
         pooling_softmax: bool = False,
-        pooling_step_tag_id: int = -1,
+        pooling_step_tag_id: Optional[int] = None,
         pooling_returned_token_ids: Optional[List[int]] = None,
         **kwargs,
     ) -> None:

diff --git a/vllm/model_executor/layers/pooler.py b/vllm/model_executor/layers/pooler.py
@@ -33,9 +33,9 @@ class Pooler(nn.Module):
     def __init__(
         self,
         pooling_type: PoolingType,
-        normalize: bool,
+        normalize: bool = False,
         softmax: bool = False,
-        step_tag_id: int = -1,
+        step_tag_id: Optional[int] = None,
         returned_token_ids: Optional[List[int]] = None,
     ):
         super().__init__()
@@ -81,7 +81,7 @@ def forward(
             pooled_data = []
             for prompt_len, seq_data_i in zip(
                     prompt_lens, pooling_metadata.seq_data.values()):
-                if self.step_tag_id == -1:
+                if self.step_tag_id is None:
                     pooled_data.append(logits[offset:offset + prompt_len])
                 else:
                     step_idxs = torch.tensor(