uv-scripts/vllm

1# /// script2# requires-python = ">=3.10"3# dependencies = [4#     "datasets",5#     "httpx",6#     "huggingface-hub",7#     "setuptools",8#     "toolz",9#     "torch",10#     "transformers",11#     "vllm>=0.11.0",12# ]13# ///14"""15Batch text classification using vLLM for efficient GPU inference.16 17This script loads a dataset from Hugging Face Hub, performs classification using18a BERT-style model via vLLM, and saves the results back to the Hub with predicted19labels and confidence scores.20 21Example usage:22    # Local execution23    uv run classify-dataset.py \\24        davanstrien/ModernBERT-base-is-new-arxiv-dataset \\25        username/input-dataset \\26        username/output-dataset \\27        --inference-column text \\28        --batch-size 1000029 30    # HF Jobs execution (see script output for full command)31    hfjobs run --flavor l4x1 ...32"""33 34import argparse35import logging36import os37import sys38from typing import Optional39 40import httpx41import torch42import torch.nn.functional as F43import vllm44from datasets import load_dataset45from huggingface_hub import hf_hub_url, login46from toolz import concat, keymap, partition_all47from tqdm.auto import tqdm48from vllm import LLM49 50logging.basicConfig(51    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"52)53logger = logging.getLogger(__name__)54 55 56def check_gpu_availability():57    """Check if CUDA is available and log GPU information."""58    if not torch.cuda.is_available():59        logger.error("CUDA is not available. This script requires a GPU.")60        logger.error(61            "Please run on a machine with NVIDIA GPU or use HF Jobs with GPU flavor."62        )63        sys.exit(1)64 65    gpu_name = torch.cuda.get_device_name(0)66    gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**367    logger.info(f"GPU detected: {gpu_name} with {gpu_memory:.1f} GB memory")68    logger.info(f"vLLM version: {vllm.__version__}")69 70 71def get_model_id2label(hub_model_id: str) -> Optional[dict[int, str]]:72    """Extract label mapping from model's config.json on Hugging Face Hub."""73    try:74        response = httpx.get(75            hf_hub_url(hub_model_id, filename="config.json"), follow_redirects=True76        )77        if response.status_code != 200:78            logger.warning(f"Could not fetch config.json for {hub_model_id}")79            return None80 81        data = response.json()82        id2label = data.get("id2label")83 84        if id2label is None:85            logger.info("No id2label mapping found in config.json")86            return None87 88        # Convert string keys to integers89        label_map = keymap(int, id2label)90        logger.info(f"Found label mapping: {label_map}")91        return label_map92 93    except Exception as e:94        logger.warning(f"Failed to parse config.json: {e}")95        return None96 97 98def get_top_label(output, label_map: Optional[dict[int, str]] = None):99    """100    Extract the top predicted label and confidence score from vLLM output.101 102    Args:103        output: vLLM ClassificationRequestOutput104        label_map: Optional mapping from label indices to label names105 106    Returns:107        Tuple of (label, confidence_score)108    """109    logits = torch.tensor(output.outputs.probs)110    probs = F.softmax(logits, dim=0)111    top_idx = torch.argmax(probs).item()112    top_prob = probs[top_idx].item()113 114    # Use label name if mapping available, otherwise use index115    label = label_map.get(top_idx, str(top_idx)) if label_map else str(top_idx)116    return label, top_prob117 118 119def main(120    hub_model_id: str,121    src_dataset_hub_id: str,122    output_dataset_hub_id: str,123    inference_column: str = "text",124    batch_size: int = 10_000,125    hf_token: Optional[str] = None,126):127    """128    Main classification pipeline.129 130    Args:131        hub_model_id: Hugging Face model ID for classification132        src_dataset_hub_id: Input dataset on Hugging Face Hub133        output_dataset_hub_id: Where to save results on Hugging Face Hub134        inference_column: Column name containing text to classify135        batch_size: Number of examples to process at once136        hf_token: Hugging Face authentication token137    """138    # GPU check139    check_gpu_availability()140 141    # Authentication142    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")143    if HF_TOKEN:144        login(token=HF_TOKEN)145    else:146        logger.error(147            "HF_TOKEN is required. Set via --hf-token or HF_TOKEN environment variable."148        )149        sys.exit(1)150 151    # vLLM auto-detects the sequence-classification runner from the model152    # architecture (*ForSequenceClassification); the task= arg was removed in 0.12.0.153    logger.info(f"Loading model: {hub_model_id}")154    llm = LLM(model=hub_model_id)155 156    # Get label mapping if available157    id2label = get_model_id2label(hub_model_id)158 159    # Load dataset160    logger.info(f"Loading dataset: {src_dataset_hub_id}")161    dataset = load_dataset(src_dataset_hub_id, split="train")162    total_examples = len(dataset)163    logger.info(f"Dataset loaded with {total_examples:,} examples")164 165    # Extract text column166    if inference_column not in dataset.column_names:167        logger.error(168            f"Column '{inference_column}' not found. Available columns: {dataset.column_names}"169        )170        sys.exit(1)171 172    prompts = dataset[inference_column]173 174    # Process in batches175    logger.info(f"Starting classification with batch size {batch_size:,}")176    all_results = []177 178    for batch in tqdm(179        list(partition_all(batch_size, prompts)),180        desc="Processing batches",181        unit="batch",182    ):183        batch_results = llm.classify(batch)184        all_results.append(batch_results)185 186    # Flatten results187    outputs = list(concat(all_results))188 189    # Extract labels and probabilities190    logger.info("Extracting predictions...")191    labels_and_probs = [get_top_label(output, id2label) for output in outputs]192 193    # Add results to dataset194    dataset = dataset.add_column("label", [label for label, _ in labels_and_probs])195    dataset = dataset.add_column("prob", [prob for _, prob in labels_and_probs])196 197    # Push to hub198    logger.info(f"Pushing results to: {output_dataset_hub_id}")199    dataset.push_to_hub(output_dataset_hub_id, token=HF_TOKEN)200    logger.info("✅ Classification complete!")201 202 203if __name__ == "__main__":204    if len(sys.argv) > 1:205        parser = argparse.ArgumentParser(206            description="Classify text data using vLLM and save results to Hugging Face Hub",207            formatter_class=argparse.RawDescriptionHelpFormatter,208            epilog="""209Examples:210  # Basic usage211  uv run classify-dataset.py model/name input-dataset output-dataset212  213  # With custom column and batch size214  uv run classify-dataset.py model/name input-dataset output-dataset \\215    --inference-column prompt \\216    --batch-size 50000217  218  # Using environment variable for token219  HF_TOKEN=hf_xxx uv run classify-dataset.py model/name input-dataset output-dataset220            """,221        )222 223        parser.add_argument(224            "hub_model_id",225            help="Hugging Face model ID for classification (e.g., bert-base-uncased)",226        )227        parser.add_argument(228            "src_dataset_hub_id",229            help="Input dataset on Hugging Face Hub (e.g., username/dataset-name)",230        )231        parser.add_argument(232            "output_dataset_hub_id", help="Output dataset name on Hugging Face Hub"233        )234        parser.add_argument(235            "--inference-column",236            type=str,237            default="text",238            help="Column containing text to classify (default: text)",239        )240        parser.add_argument(241            "--batch-size",242            type=int,243            default=10_000,244            help="Batch size for inference (default: 10,000)",245        )246        parser.add_argument(247            "--hf-token",248            type=str,249            help="Hugging Face token (can also use HF_TOKEN env var)",250        )251 252        args = parser.parse_args()253 254        main(255            hub_model_id=args.hub_model_id,256            src_dataset_hub_id=args.src_dataset_hub_id,257            output_dataset_hub_id=args.output_dataset_hub_id,258            inference_column=args.inference_column,259            batch_size=args.batch_size,260            hf_token=args.hf_token,261        )262    else:263        # Show HF Jobs example when run without arguments264        print("""265vLLM Classification Script266=========================267 268This script requires arguments. For usage information:269    uv run classify-dataset.py --help270 271Example HF Jobs command:272    hfjobs run \\273        --flavor l4x1 \\274        --secret HF_TOKEN=$(python -c "from huggingface_hub import HfFolder; print(HfFolder.get_token())") \\275        vllm/vllm-openai:latest \\276        /bin/bash -c '277            uv run https://huggingface.co/datasets/uv-scripts/vllm/resolve/main/classify-dataset.py \\278                davanstrien/ModernBERT-base-is-new-arxiv-dataset \\279                username/input-dataset \\280                username/output-dataset \\281                --inference-column text \\282                --batch-size 100000283        ' \\284        --project vllm-classify \\285        --name my-classification-job286        """)287