Added extra_repr() to pooler classes to improve debuggability (#44805)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

vllm/model_executor/layers/pooler/activations.py

@@ -128,6 +128,9 @@ class PoolerClassify(PoolerActivation):
 
         self.num_labels = num_labels
 
+    def extra_repr(self) -> str:
+        return f"num_labels={self.num_labels}"
+
     def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
         num_labels = self.num_labels
         if num_labels is None:
@@ -145,5 +148,11 @@ class LambdaPoolerActivation(PoolerActivation):
 
         self.fn = fn
 
+    def extra_repr(self) -> str:
+        name = getattr(self.fn, "__name__", None)
+        if name is None:
+            name = self.fn.__class__.__name__
+        return f"fn={name}"
+
     def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
         return self.fn(pooled_data)

vllm/model_executor/layers/pooler/seqwise/heads.py

@@ -43,6 +43,16 @@ class EmbeddingPoolerHead(SequencePoolerHead):
         self.head_dtype = head_dtype
         self.activation = activation
 
+    def extra_repr(self) -> str:
+        attrs = []
+        if self.head_dtype is not None:
+            attrs.append(f"head_dtype={self.head_dtype}")
+        if self.projector is not None:
+            attrs.append("projector=True")
+        if self.activation is not None:
+            attrs.append(f"activation={self.activation.__class__.__name__}")
+        return ", ".join(attrs)
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return {"embed"}
 
@@ -124,6 +134,20 @@ class ClassifierPoolerHead(SequencePoolerHead):
         self.head_dtype = head_dtype
         self.activation = activation
 
+    def extra_repr(self) -> str:
+        attrs = []
+        if self.head_dtype is not None:
+            attrs.append(f"head_dtype={self.head_dtype}")
+        if self.classifier is not None:
+            attrs.append("classifier=True")
+        if self.logit_mean is not None:
+            attrs.append(f"logit_mean={self.logit_mean}")
+        if self.logit_sigma is not None:
+            attrs.append(f"logit_sigma={self.logit_sigma}")
+        if self.activation is not None:
+            attrs.append(f"activation={self.activation.__class__.__name__}")
+        return ", ".join(attrs)
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return {"classify"}
 

vllm/model_executor/layers/pooler/seqwise/poolers.py

@@ -61,6 +61,12 @@ class SequencePooler(Pooler):
         self.pooling = pooling
         self.head = head
 
+    def extra_repr(self) -> str:
+        return (
+            f"pooling={self.pooling.__class__.__name__}, "
+            f"head={self.head.__class__.__name__}"
+        )
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         tasks = set(POOLING_TASKS)
 

vllm/model_executor/layers/pooler/special.py

@@ -164,6 +164,9 @@ class BOSEOSFilter(Pooler):
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
 
+    def extra_repr(self) -> str:
+        return f"bos_token_id={self.bos_token_id}, eos_token_id={self.eos_token_id}"
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return self.pooler.get_supported_tasks()
 

vllm/model_executor/layers/pooler/tokwise/heads.py

@@ -58,6 +58,16 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
         self.projector = projector
         self.activation = activation
 
+    def extra_repr(self) -> str:
+        attrs = []
+        if self.head_dtype is not None:
+            attrs.append(f"head_dtype={self.head_dtype}")
+        if self.projector is not None:
+            attrs.append("projector=True")
+        if self.activation is not None:
+            attrs.append(f"activation={self.activation.__class__.__name__}")
+        return ", ".join(attrs)
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return {"token_embed"}
 
@@ -110,6 +120,20 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
         self.head_dtype = head_dtype
         self.activation = activation
 
+    def extra_repr(self) -> str:
+        attrs = []
+        if self.head_dtype is not None:
+            attrs.append(f"head_dtype={self.head_dtype}")
+        if self.classifier is not None:
+            attrs.append("classifier=True")
+        if self.logit_mean is not None:
+            attrs.append(f"logit_mean={self.logit_mean}")
+        if self.logit_sigma is not None:
+            attrs.append(f"logit_sigma={self.logit_sigma}")
+        if self.activation is not None:
+            attrs.append(f"activation={self.activation.__class__.__name__}")
+        return ", ".join(attrs)
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         return {"token_classify"}
 

vllm/model_executor/layers/pooler/tokwise/methods.py

@@ -41,6 +41,9 @@ class AllPool(TokenPoolingMethod):
 
         self.enable_chunked_prefill = scheduler_config.enable_chunked_prefill
 
+    def extra_repr(self) -> str:
+        return f"enable_chunked_prefill={self.enable_chunked_prefill}"
+
     def forward(
         self,
         hidden_states: torch.Tensor,

vllm/model_executor/layers/pooler/tokwise/poolers.py

@@ -65,6 +65,10 @@ class TokenPooler(Pooler):
         self.pooling = pooling
         self.head = head
 
+    def extra_repr(self) -> str:
+        head_name = self.head.__class__.__name__ if self.head is not None else None
+        return f"pooling={self.pooling.__class__.__name__}, head={head_name}"
+
     def get_supported_tasks(self) -> Set[PoolingTask]:
         tasks = set(POOLING_TASKS)