[Usage]: How to get the logits from classification models.

[kxzxvbk]

Your current environment

pip install vllm

How would you like to use vllm

Hi! When I use vllm for classification, I find that the returned output is something like: ClassificationOutput(num_classes=5), which only contains the probabilities.
Now, I want to directly get the logits (i.e. vectors before softmax).
I think it sounds easy, but I cannot find a good way. Anyone can help me?

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[noooop]

try

override_pooler_config = {"normalize": False}

logits should be "softmax": False instead of "normalize": False

[kxzxvbk]

I'm using the offline mode, and here is my code:

# SPDX-License-Identifier: Apache-2.0

from argparse import Namespace
from easydict import EasyDict

from vllm import LLM, EngineArgs
from vllm.utils import FlexibleArgumentParser


def parse_args():
    parser = FlexibleArgumentParser()
    parser = EngineArgs.add_cli_args(parser)
    # Set example specific arguments
    parser.set_defaults(model="./models/qwen_8b_reward_model",
                        task="classify",
                        override_pooler_config=EasyDict({"pooling_type": "LAST", "normalize": False, "softmax": None, "step_tag_id": None, "returned_token_ids": None}))
    return parser.parse_args()


def main(args: Namespace):
    # Sample prompts.
    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]

    # Create an LLM.
    # You should pass task="classify" for classification models
    model = LLM(**vars(args))

    # Generate logits. The output is a list of ClassificationRequestOutputs.
    outputs = model.classify(prompts)

    # Print the outputs.
    print("\nGenerated Outputs:\n" + "-" * 60)
    for prompt, output in zip(prompts, outputs):
        probs = output.outputs
        print(f"Prompt: {prompt!r} \n"
              f"Class Probabilities: {probs.__dict__}")
        print("-" * 60)


if __name__ == "__main__":
    args = parse_args()
    main(args)

The output is:

Generated Outputs:
------------------------------------------------------------
Prompt: 'Hello, my name is' 
Class Probabilities: {'probs': xxx}
------------------------------------------------------------
Prompt: 'The president of the United States is' 
Class Probabilities: {'probs': xxx}
------------------------------------------------------------
Prompt: 'The capital of France is' 
Class Probabilities: {'probs': xxx}
------------------------------------------------------------
Prompt: 'The future of AI is' 
Class Probabilities: {'probs': xxx}

Seems that it doesn't work.

[kxzxvbk]

Oh! I find the error.
This should be modified as: {"pooling_type": "LAST", "normalize": False, "softmax": False, "step_tag_id": None, "returned_token_ids": None}
And the result will be correct. Thanks for your help!