[Misc]Add customized information for models (#4132)

2024-05-01 12:18:14 +08:00
parent c3845d82dc
commit d6f4bd7cdd
9 changed files with 87 additions and 0 deletions
--- a/vllm/model_executor/layers/logits_processor.py
+++ b/vllm/model_executor/layers/logits_processor.py
@@ -70,6 +70,12 @@ class LogitsProcessor(nn.Module):
            logits = logits[:, :self.org_vocab_size]
        return logits

+    def extra_repr(self) -> str:
+        s = f"vocab_size={self.vocab_size}"
+        s += f", forg_vocab_size={self.org_vocab_size}"
+        s += f", scale={self.scale}, logits_as_input={self.logits_as_input}"
+        return s
+

 def _prune_hidden_states(
    hidden_states: torch.Tensor,