wip

Signed-off-by: Flavia Beo <[email protected]>
flaviabeo · Oct 21, 2024 · ba2c001 · ba2c001
1 parent 1cb5036
commit ba2c001
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/vllm/model_executor/layers/pooler.py b/vllm/model_executor/layers/pooler.py
@@ -37,7 +37,7 @@ def get_pooling_type(self, pooling_type_name: str) -> PoolingType:
         return PoolingType(next((value for key, 
                              value in pooling_types if key.lower() 
                              in pooling_type_name), 
-                             None))
+                             2))
 
 
 class Pooler(nn.Module):

diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py
@@ -391,7 +391,8 @@ def __init__(
     ) -> None:
         super().__init__()
         self.model = BertModel(config, cache_config, quant_config)
-        print(pooling_config)
+        print(pooling_config.pooling_type)
+        print(pooling_config.normalize)
         self._pooler = Pooler(pooling_config.pooling_type, 
                               pooling_config.normalize)