Text Classification

def classify(
    self,
    prompts: Union[PromptType, Sequence[PromptType]],
    *,
    use_tqdm: Union[bool, Callable[..., tqdm]] = True,
    pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,
    lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None
) -> List[ClassificationRequestOutput]:
    """
    Classify text into predefined categories.

    Parameters:
    - prompts: Input text to classify
    - use_tqdm: Show progress bar (keyword-only)
    - pooling_params: Pooling configuration (keyword-only)
    - lora_request: LoRA adapter configuration (keyword-only)

    Returns:
    List of ClassificationRequestOutput with predictions

    Note: This method requires a model that supports classification tasks.
    The class labels are determined by the model's configuration, not passed as parameters.
    """

from vllm import LLM, PoolingParams

# Use a classification model that has predefined classes
llm = LLM(model="cardiffnlp/twitter-roberta-base-sentiment-latest")

texts = [
    "I love this product! It's amazing.",
    "This is terrible, worst purchase ever.",
    "It's okay, not great but not bad."
]

pooling_params = PoolingParams(pooling_type="MEAN")

outputs = llm.classify(texts, pooling_params=pooling_params)

for output in outputs:
    result = output.outputs
    print(f"Text: {texts[outputs.index(output)]}")
    print(f"Predicted probabilities: {result.probs}")

class ClassificationRequestOutput:
    id: str
    outputs: ClassificationOutput
    prompt_token_ids: List[int]
    finished: bool

class ClassificationOutput:
    probs: List[float]  # Class probabilities