scandeval.benchmark_modules.vllm

docs module scandeval.benchmark_modules.vllm

"""Generative models using the vLLM inference framework."""

import collections.abc as c
import importlib.util
import itertools as it
import json
import logging
import os
import random
import re
import sys
import typing as t
from functools import partial
from pathlib import Path
from time import sleep
from types import MethodType

import torch
from datasets import DatasetDict
from huggingface_hub import snapshot_download
from pydantic import conlist, create_model
from tqdm.auto import tqdm
from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizer, Trainer
from urllib3.exceptions import RequestError

from ..constants import (
    GENERATIVE_PIPELINE_TAGS,
    MAX_LOGPROBS,
    MERGE_TAGS,
    REASONING_MAX_TOKENS,
    TASK_GROUPS_USING_LOGPROBS,
    TASKS_USING_JSON,
)
from ..data_models import (
    BenchmarkConfig,
    DatasetConfig,
    GenerativeModelOutput,
    ModelConfig,
    Task,
)
from ..enums import (
    BatchingPreference,
    GenerativeType,
    InferenceBackend,
    ModelType,
    TaskGroup,
)
from ..exceptions import (
    InvalidBenchmark,
    InvalidModel,
    NeedsEnvironmentVariable,
    NeedsExtraInstalled,
)
from ..languages import get_all_languages
from ..task_utils import (
    question_answering,
    sequence_classification,
    text_to_text,
    token_classification,
)
from ..types import ExtractLabelsFunction
from ..utils import (
    clear_memory,
    create_model_cache_dir,
    get_end_of_chat_token_ids,
    log_once,
    should_prompts_be_stripped,
)
from .hf import HuggingFaceEncoderModel, get_model_repo_info, load_hf_model_config

if t.TYPE_CHECKING or importlib.util.find_spec("vllm") is not None:
    from vllm import LLM, RequestOutput, SamplingParams
    from vllm.lora.request import LoRARequest
    from vllm.sampling_params import GuidedDecodingParams

    try:
        from vllm.model_executor.parallel_utils.parallel_state import (
            destroy_model_parallel,
        )
    except ImportError:
        from vllm.distributed.parallel_state import destroy_model_parallel

if t.TYPE_CHECKING or importlib.util.find_spec("ray") is not None:
    import ray

logger = logging.getLogger("scandeval")


class VLLMModel(HuggingFaceEncoderModel):
docs
    """A generative model using the vLLM inference framework."""

    fresh_model = False
    batching_preference = BatchingPreference.ALL_AT_ONCE
    high_priority = True

    def __init__(
        self,
        model_config: ModelConfig,
        dataset_config: DatasetConfig,
        benchmark_config: BenchmarkConfig,
    ) -> None:
        """Initialise the vLLM model.

        Args:
            model_config:
                The model configuration.
            dataset_config:
                The dataset configuration.
            benchmark_config:
                The benchmark configuration.
        """
        if (
            importlib.util.find_spec("vllm") is None
            or importlib.util.find_spec("ray") is None
        ):
            raise NeedsExtraInstalled(extra="generative")

        output_scores = dataset_config.task.task_group in TASK_GROUPS_USING_LOGPROBS
        model, tokenizer = load_model_and_tokenizer(
            model_config=model_config,
            benchmark_config=benchmark_config,
            output_scores=output_scores,
        )
        self._model: LLM = model
        self._tokenizer: PreTrainedTokenizer = tokenizer
        self.end_of_reasoning_token_id = get_end_of_reasoning_token_id(
            model=self._model, tokenizer=self._tokenizer
        )

        # We specify `HuggingFaceEncoderModel` here instead of `VLLMModel`, as we want
        # to call the `__init__` method of the `BenchmarkModule` class.
        super(HuggingFaceEncoderModel, self).__init__(
            model_config=model_config,
            dataset_config=dataset_config,
            benchmark_config=benchmark_config,
        )

        self.buffer["output_scores"] = output_scores
        self.buffer["instruction_model"] = self._tokenizer.chat_template is not None
        if self.model_config.adapter_base_model_id is not None:
            adapter_path = snapshot_download(
                repo_id=self.model_config.model_id,
                cache_dir=Path(self.model_config.model_cache_dir),
            )
            self.buffer["lora_request"] = LoRARequest(
                lora_name="adapter", lora_int_id=1, lora_path=adapter_path
            )

    @property
    def generative_type(self) -> GenerativeType | None:
docs
        """Get the generative type of the model.

        Returns:
            The generative type of the model, or None if it has not been set yet.
        """
        if not hasattr(self, "_tokenizer"):
            return None
        elif self.end_of_reasoning_token_id is not None:
            return GenerativeType.REASONING
        elif self._tokenizer.chat_template is not None:
            return GenerativeType.INSTRUCTION_TUNED
        else:
            return GenerativeType.BASE

    @property
    def extract_labels_from_generation(self) -> ExtractLabelsFunction:
docs
        """The function used to extract the labels from the generated output.

        Returns:
            The function used to extract the labels from the generated output.
        """
        match self.dataset_config.task.task_group:
            case (
                TaskGroup.SEQUENCE_CLASSIFICATION
                | TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION
            ):
                return partial(
                    sequence_classification.extract_labels_from_generation,
                    dataset_config=self.dataset_config,
                )
            case TaskGroup.TEXT_TO_TEXT:
                return text_to_text.extract_labels_from_generation
            case TaskGroup.TOKEN_CLASSIFICATION:
                return partial(
                    token_classification.extract_labels_from_generation,
                    dataset_config=self.dataset_config,
                )
            case TaskGroup.QUESTION_ANSWERING:
                return question_answering.extract_labels_from_generation
            case _:
                raise NotImplementedError(
                    f"Unsupported task group: {self.dataset_config.task.task_group}."
                )

    def prepare_dataset(
docs
        self, dataset: DatasetDict, task: Task, itr_idx: int
    ) -> DatasetDict:
        """Prepare the dataset for the model.

        This includes things like tokenisation.

        Args:
            dataset:
                The dataset to prepare.
            task:
                The task to prepare the dataset for.
            itr_idx:
                The index of the dataset in the iterator.

        Returns:
            The prepared dataset.
        """
        if task.task_group == TaskGroup.QUESTION_ANSWERING:
            dataset = dataset.map(
                lambda examples: dict(
                    label=[
                        dict(
                            id=id,
                            answers=dict(
                                answer_start=answer_dct["answer_start"],
                                text=[
                                    answer_text.lower()
                                    for answer_text in answer_dct["text"]
                                ],
                            ),
                        )
                        for id, answer_dct in zip(examples["id"], examples["answers"])
                    ]
                ),
                batched=True,
                load_from_cache_file=False,
                keep_in_memory=True,
            )

        if self.benchmark_config.few_shot:
            few_shot_examples = self._extract_few_shot_examples(
                dataset=dataset, task=task, itr_idx=itr_idx
            )
        else:
            few_shot_examples = list()

        dataset["test"] = dataset["test"].map(
            partial(self._apply_prompt, few_shot_examples=few_shot_examples, task=task),
            batched=True,
            load_from_cache_file=False,
            keep_in_memory=True,
        )

        return dataset

    def generate(self, inputs: dict) -> GenerativeModelOutput:
docs
        """Generate outputs from the model.

        Args:
            inputs:
                A batch of inputs to pass through the model.

        Returns:
            The generated model outputs.
        """
        # Define which tokens to use as stopping criteria. We want to use the padding
        # token, end-of-sentence token, and a double newline if the model isn't
        # instruction tuned (since these separate the few-shot examples in the input in
        # this case)
        stop_tokens: list[str] = list()
        if self.buffer["instruction_model"] is False:
            stop_tokens.append("\n\n")
        if self._tokenizer.pad_token_id is not None:
            stop_tokens.append(self._tokenizer.pad_token)
        if self._tokenizer.eos_token_id is not None:
            stop_tokens.append(self._tokenizer.eos_token)
            if self._tokenizer.pad_token_id is None:
                self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
                self._tokenizer.pad_token = self._tokenizer.eos_token
        if (
            self._tokenizer.bos_token_id is not None
            and self._tokenizer.pad_token_id is None
        ):
            self._tokenizer.pad_token_id = self._tokenizer.bos_token_id
            self._tokenizer.pad_token = self._tokenizer.bos_token
        elif (
            self._tokenizer.eos_token_id is not None
            and self._tokenizer.pad_token_id is None
        ):
            self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
            self._tokenizer.pad_token = self._tokenizer.eos_token
        elif self._tokenizer.pad_token_id is None:
            pad_token_candidates = ["<pad>", "[pad]", "<|endoftext|>", "<|im_end|>"]
            pad_token_candidates.extend([c.upper() for c in pad_token_candidates])
            for candidate in pad_token_candidates:
                if candidate in self._tokenizer.get_vocab():
                    pad_token_id = self._tokenizer.get_vocab()[candidate]
                    self._tokenizer.pad_token = candidate
                    self._tokenizer.pad_token_id = pad_token_id
                    break
            else:
                raise InvalidModel(
                    "Could not find a suitable token to use as a padding token, since "
                    "the model does not have a BOS, EOS, or padding token, and does "
                    f"not have any of the following tokens in its vocabulary: "
                    f"{pad_token_candidates}."
                )

        assert self._tokenizer.pad_token_id is not None

        # Add end of chat token as a stopping token, if it exists
        end_of_chat_token_ids = get_end_of_chat_token_ids(tokenizer=self._tokenizer)
        if end_of_chat_token_ids is not None:
            end_of_chat_token = self._tokenizer.decode(end_of_chat_token_ids).strip()
            if end_of_chat_token:
                stop_tokens.append(end_of_chat_token)

        if self.dataset_config.task in TASKS_USING_JSON:
            ner_tag_names = list(self.dataset_config.prompt_label_mapping.values())
            keys_and_their_types: dict[str, t.Any] = {
                tag_name: (conlist(str, max_length=5), ...)
                for tag_name in ner_tag_names
            }
            pydantic_class = create_model("AnswerFormat", **keys_and_their_types)
            schema = pydantic_class.model_json_schema()
            guided_decoding = GuidedDecodingParams(
                json=schema, backend="outlines", whitespace_pattern=r" ?"
            )
        else:
            guided_decoding = None

        # Define the parameters used for vLLM generation
        max_tokens: int = (
            REASONING_MAX_TOKENS
            if self.generative_type == GenerativeType.REASONING
            else self.dataset_config.max_generated_tokens
        )
        sampling_params = SamplingParams(
            max_tokens=max_tokens,
            logprobs=MAX_LOGPROBS if self.buffer["output_scores"] else None,
            temperature=0.0,
            stop=[stop_token for stop_token in stop_tokens if stop_token],
            guided_decoding=guided_decoding,
        )

        # If any of the prompts are empty then we need to replace them with a BOS token
        # so that the vLLM model can generate from them
        prompts: list[str] = inputs["text"]
        if any(len(prompt) == 0 for prompt in prompts):
            logger.debug("Found empty prompts, replacing with BOS token.")
            prompts = [
                prompt if len(prompt) > 0 else str(self._tokenizer.bos_token)
                for prompt in prompts
            ]

        # Strip the prompts if the model's tokeniser requires it
        labels_to_be_generated = list(self.dataset_config.prompt_label_mapping.values())
        if len(labels_to_be_generated) == 0:
            labels_to_be_generated = ["negative", "positive"]
        if not self.buffer.get(
            "instruction_model", False
        ) and should_prompts_be_stripped(
            labels_to_be_generated=labels_to_be_generated, tokenizer=self._tokenizer
        ):
            log_once(message="Stripping prompts.", level=logging.DEBUG)
            prompts = [prompt.strip() for prompt in prompts]

        # Generate sequences using vLLM
        input_is_a_test = len(prompts) == 1 and len(set(prompts[0])) == 1
        raw_outputs = self._model.generate(
            prompts=prompts,
            sampling_params=sampling_params,
            use_tqdm=(not input_is_a_test),
            lora_request=self.buffer.get("lora_request"),
        )
        completion_ids: list[list[int]] = [
            output.outputs[0].token_ids for output in raw_outputs
        ]
        if self.end_of_reasoning_token_id in completion_ids[0]:
            completion_ids = [
                token_ids[token_ids.index(self.end_of_reasoning_token_id) + 2 :]
                if self.end_of_reasoning_token_id in token_ids
                else token_ids
                for token_ids in completion_ids
            ]
        completions = self._tokenizer.batch_decode(
            sequences=[
                torch.LongTensor(completion_id) for completion_id in completion_ids
            ],
            skip_special_tokens=True,
        )
        completions = [completion.strip() for completion in completions]

        # Add logprobs scores to the output
        if self.buffer["output_scores"]:
            scores: list[list[list[tuple[str, float]]]] = [
                [
                    [
                        (obj.decoded_token, obj.logprob)
                        for obj in token_logprobs_dict.values()
                    ]
                    for token_logprobs_dict in raw_output.outputs[0].logprobs
                ]
                for raw_output in raw_outputs
            ]
            scores = [
                score_list[
                    raw_output.outputs[0].token_ids.index(
                        self.end_of_reasoning_token_id
                    )
                    + 2 :
                ]
                if self.end_of_reasoning_token_id in raw_output.outputs[0].token_ids
                else score_list
                for raw_output, score_list in zip(raw_outputs, scores)
            ]
            output = GenerativeModelOutput(sequences=completions, scores=scores)
        else:
            output = GenerativeModelOutput(sequences=completions)

        return output

    @classmethod
    def model_exists(
docs
        cls, model_id: str, benchmark_config: BenchmarkConfig
    ) -> bool | NeedsExtraInstalled | NeedsEnvironmentVariable:
        """Check if a model exists.

        Args:
            model_id:
                The model ID.
            benchmark_config:
                The benchmark configuration.

        Returns:
            Whether the model exists, or an error describing why we cannot check
            whether the model exists.
        """
        using_api = (
            benchmark_config.api_base is not None
            or benchmark_config.api_version is not None
        )
        if using_api:
            return False

        model_id, revision = (
            model_id.split("@") if "@" in model_id else (model_id, "main")
        )
        model_info = get_model_repo_info(
            model_id=model_id, revision=revision, benchmark_config=benchmark_config
        )
        return (
            model_info is not None
            and model_info.pipeline_tag in GENERATIVE_PIPELINE_TAGS
        )

    @classmethod
    def get_model_config(
docs
        cls, model_id: str, benchmark_config: BenchmarkConfig
    ) -> ModelConfig:
        """Fetch the model configuration.

        Args:
            model_id:
                The model ID.
            benchmark_config:
                The benchmark configuration.

        Returns:
            The model configuration.
        """
        model_id, revision = (
            model_id.split("@") if "@" in model_id else (model_id, "main")
        )
        model_info = get_model_repo_info(
            model_id=model_id, revision=revision, benchmark_config=benchmark_config
        )
        if model_info is None:
            raise InvalidModel(f"The model {model_id!r} could not be found.")

        language_mapping = get_all_languages()
        language_codes = list(language_mapping.keys())

        model_config = ModelConfig(
            model_id=model_id,
            revision=revision,
            task=model_info.pipeline_tag,
            languages=[
                language_mapping[tag]
                for tag in model_info.tags
                if tag in language_codes
            ],
            merge=any(tag in model_info.tags for tag in MERGE_TAGS),
            inference_backend=InferenceBackend.VLLM,
            model_type=ModelType.GENERATIVE,
            fresh=False,
            model_cache_dir=create_model_cache_dir(
                cache_dir=benchmark_config.cache_dir, model_id=model_id
            ),
            adapter_base_model_id=model_info.adapter_base_model_id,
        )

        return model_config

    def _extract_few_shot_examples(
        self, dataset: DatasetDict, task: Task, itr_idx: int
    ) -> list[dict[str, t.Any]]:
        """Extract few-shot examples from a dataset.

        This will always extract the examples from the training split.

        We ensure that the few-shot examples are unique by picking them one at a time.

        Args:
            dataset:
                The dataset to extract the few-shot examples from.
            task:
                The task that is being benchmarked.
            itr_idx:
                The index of the dataset in the iterator.

        Returns:
            The few-shot examples.
        """
        random_seed = 4242 + itr_idx
        num_few_shots = self.dataset_config.num_few_shot_examples
        few_shot_examples: list[dict[str, t.Any]] = list()
        shuffled_train = dataset["train"].shuffle(seed=random_seed)

        match task.task_group:
            case (
                TaskGroup.SEQUENCE_CLASSIFICATION
                | TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION
            ):
                labels = it.cycle(self.dataset_config.labels)
                while (
                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
                ):
                    label = next(labels)
                    possible_examples = shuffled_train.filter(
                        lambda x: x["label"].lower() == label.lower()
                    )
                    if len(possible_examples) == 0:
                        continue
                    example = possible_examples.select(range(1))[0]
                    few_shot_examples.append(example)
                    shuffled_train = shuffled_train.filter(
                        lambda x: x["text"] != example["text"]
                    )

            case TaskGroup.TEXT_TO_TEXT:
                while (
                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
                ):
                    example = shuffled_train.select(range(1))[0]
                    few_shot_examples.append(example)
                    shuffled_train = shuffled_train.filter(
                        lambda x: x["text"] != example["text"]
                    )

            case TaskGroup.TOKEN_CLASSIFICATION:
                labels = it.cycle(
                    [
                        label.lower()
                        for label in self.dataset_config.labels
                        if label.lower().startswith("b-")
                    ]
                )
                while (
                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
                ):
                    label = next(labels)
                    possible_examples = shuffled_train.filter(
                        lambda x: label in [tag.lower() for tag in x["labels"]]
                    )
                    if len(possible_examples) == 0:
                        continue
                    example = possible_examples.select(range(1))[0]
                    few_shot_examples.append(example)
                    shuffled_train = shuffled_train.filter(
                        lambda x: x["tokens"] != example["tokens"]
                    )

            case TaskGroup.QUESTION_ANSWERING:
                # Locate the maximum number of tokens that constitutes a short example
                for max_num_tokens in [512, 1024, 2048, 4096, 8192]:
                    train_with_short_examples = dataset["train"].filter(
                        lambda example: len(example["context"]) < max_num_tokens
                    )
                    num_short_examples = len(train_with_short_examples)
                    if num_short_examples >= self.dataset_config.num_few_shot_examples:
                        break
                else:
                    raise InvalidBenchmark(
                        "Could not find enough short examples for few-shot learning."
                    )

                shuffled_train = train_with_short_examples.shuffle(seed=random_seed)
                while (
                    len(few_shot_examples) < num_few_shots and len(shuffled_train) > 0
                ):
                    example = shuffled_train.select(range(1))[0]
                    few_shot_examples.append(example)
                    shuffled_train = shuffled_train.filter(
                        lambda x: x["context"] != example["context"]
                    )

            case _:
                raise NotImplementedError(f"Unsupported task group: {task.task_group}.")

        random.seed(random_seed)
        random.shuffle(few_shot_examples)
        return few_shot_examples

    def _apply_prompt(
        self,
        examples: dict[str, t.Any],
        few_shot_examples: list[dict[str, t.Any]],
        task: Task,
    ) -> dict[str, t.Any]:
        """Apply prompt template to an example, potentially with few-shot examples.

        Args:
            examples:
                The examples to apply the few-shot examples to.
            few_shot_examples:
                The few-shot examples to apply.
            task:
                The task that is being benchmarked.

        Returns:
            The example with the few-shot examples applied.
        """

        def create_prompt(**kwargs) -> tuple[str, str]:
            """Create a prompt from the given keyword arguments.

            Args:
                kwargs:
                    The keyword arguments to use in the prompt.

            Returns:
                A pair (prompt, label), where "label" is an empty string if the model is
                not instruction tuned (as in this case it is included in the prompt).
            """
            label_key = "label" if "label" in kwargs else "target_text"
            label = kwargs.pop(label_key)
            assert (
                label is not None
            ), f"Found a None label for the prompt: {kwargs}. This should not happen."
            label_mapping = self.dataset_config.prompt_label_mapping
            label = label_mapping.get(label, label)
            if self.buffer["instruction_model"]:
                prompt = self.dataset_config.instruction_prompt.format(**kwargs)
                return prompt, label
            else:
                kwargs[label_key] = label
                return self.dataset_config.prompt_template.format(**kwargs), ""

        match task.task_group:
            case (
                TaskGroup.SEQUENCE_CLASSIFICATION
                | TaskGroup.MULTIPLE_CHOICE_CLASSIFICATION
            ):
                few_shot_sections = [
                    create_prompt(
                        text=example["text"].replace("\n", " ").strip(),
                        label=example["label"].replace("\n", " ").strip(),
                    )
                    for example in few_shot_examples
                ]
                new_sections = [
                    create_prompt(text=text.replace("\n", " ").strip(), label="")
                    for text in examples["text"]
                ]

            case TaskGroup.TEXT_TO_TEXT:
                few_shot_sections = [
                    create_prompt(
                        text=example["text"].replace("\n", " ").strip(),
                        target_text=example["target_text"].replace("\n", " ").strip(),
                    )
                    for example in few_shot_examples
                ]
                new_sections = [
                    create_prompt(text=text.replace("\n", " ").strip(), target_text="")
                    for text in examples["text"]
                ]

            case TaskGroup.TOKEN_CLASSIFICATION:

                def create_label(example: dict) -> str:
                    prompt_labels = self.dataset_config.prompt_label_mapping.values()
                    labels: dict[str, list[str]] = {
                        prompt_label: list() for prompt_label in prompt_labels
                    }
                    for token, label in zip(example["tokens"], example["labels"]):
                        label = label.lower()
                        if label == "o":
                            continue
                        prompt_label = self.dataset_config.prompt_label_mapping[label]
                        if label.startswith("b-"):
                            labels[prompt_label].append(token)
                        elif label.startswith("i-"):
                            labels[prompt_label][-1] += " " + token
                    return json.dumps(labels, ensure_ascii=False)

                few_shot_sections = [
                    create_prompt(
                        text=" ".join(example["tokens"]).replace("\n", " ").strip(),
                        label=create_label(example=example),
                    )
                    for example in few_shot_examples
                ]
                new_sections = [
                    create_prompt(
                        text=" ".join(tokens).replace("\n", " ").strip(), label=""
                    )
                    for tokens in examples["tokens"]
                ]

            case TaskGroup.QUESTION_ANSWERING:
                few_shot_sections = [
                    create_prompt(
                        text=example["context"].replace("\n", " ").strip(),
                        question=example["question"].replace("\n", " ").strip(),
                        label=example["answers"]["text"][0].replace("\n", " "),
                    )
                    for example in few_shot_examples
                ]
                new_sections = [
                    create_prompt(
                        text=context.replace("\n", " ").strip(),
                        question=question.replace("\n", " ").strip(),
                        label="",
                    )
                    for context, question in zip(
                        examples["context"], examples["question"]
                    )
                ]

            case _:
                raise NotImplementedError(f"Unsupported task group: {task.task_group}.")

        if self.buffer["instruction_model"]:
            few_shot_messages = [
                dict(role=role, content=content)
                for prompt, label in few_shot_sections
                for role, content in [("user", prompt), ("assistant", label)]
            ]

            messages_list = [
                few_shot_messages + [dict(role="user", content=prompt)]
                for prompt, _ in new_sections
            ]

            # Pick the chat template that matches the language of the dataset, if such a
            # template exists
            chat_template: str | None = None
            if isinstance(self._tokenizer.chat_template, dict):
                language_codes = [
                    language.code for language in self.dataset_config.languages
                ]
                for name, candidate_template in self._tokenizer.chat_template.items():
                    if name.lower() in language_codes:
                        chat_template = candidate_template
                        log_once(
                            f"Using the {name!r} chat template for the tokenizer.",
                            level=logging.DEBUG,
                        )
                        break

            texts = [
                self._tokenizer.apply_chat_template(
                    conversation=messages,
                    tokenize=False,
                    add_generation_prompt=True,
                    chat_template=chat_template,
                )
                for messages in messages_list
            ]

            examples["text"] = texts

        else:
            prompt_prefix = ""
            if self.dataset_config.prompt_prefix:
                prompt_prefix = self.dataset_config.prompt_prefix + "\n\n"

            few_shot_prompt = "\n\n".join([prompt for prompt, _ in few_shot_sections])
            if few_shot_prompt:
                few_shot_prompt += "\n\n"

            examples["text"] = [
                prompt_prefix + few_shot_prompt + new_prompt
                for new_prompt, _ in new_sections
            ]

        return examples

    @property
docs
    def data_collator(self) -> c.Callable[[list[t.Any]], dict[str, t.Any]]:
        """The data collator used to prepare samples during finetuning.

        Returns:
            The data collator.
        """
        raise NotImplementedError(
            "The `data_collator` property has not been implemented for vLLM models."
        )

    @property
    def trainer_class(self) -> t.Type["Trainer"]:
docs
        """The Trainer class to use for finetuning.

        Returns:
            The Trainer class.
        """
        raise NotImplementedError(
            "The `trainer_class` property has not been implemented for vLLM models."
        )


def load_model_and_tokenizer(
docs
    model_config: ModelConfig, benchmark_config: BenchmarkConfig, output_scores: bool
) -> "tuple[LLM, PreTrainedTokenizer]":
    """Load the model and tokenizer.

    Args:
        model_config:
            The model configuration.
        benchmark_config:
            The benchmark configuration.
        output_scores:
            Whether to output scores.

    Returns:
        The loaded model and tokenizer.
    """
    # Prefer base model ID if the model is an adapter - the adapter will be added on
    # during inference in this case
    model_id = model_config.adapter_base_model_id or model_config.model_id

    hf_model_config = load_hf_model_config(
        model_id=model_id,
        num_labels=0,
        id2label=dict(),
        label2id=dict(),
        revision=model_config.revision,
        model_cache_dir=model_config.model_cache_dir,
        api_key=benchmark_config.api_key,
        trust_remote_code=benchmark_config.trust_remote_code,
        run_with_cli=benchmark_config.run_with_cli,
    )

    quantization = None
    if hasattr(hf_model_config, "quantization_config"):
        quantization = hf_model_config.quantization_config.get("quant_method")

    # The quantised models require extra dependencies
    if quantization == "gptq" and (
        importlib.util.find_spec("auto_gptq") is None
        or importlib.util.find_spec("optimum") is None
    ):
        raise NeedsExtraInstalled(extra="quantization")
    if quantization == "awq" and importlib.util.find_spec("awq") is None:
        raise NeedsExtraInstalled(extra="quantization")

    dtype: str | torch.dtype = "auto"
    if quantization is not None and hf_model_config.torch_dtype != torch.float16:
        logger.info(
            "You are loading a quantized model with dtype "
            f"{hf_model_config.torch_dtype}, which vLLM does not support. Setting "
            "dtype to float16 instead."
        )
        dtype = torch.float16

    if model_config.adapter_base_model_id is not None:
        download_dir = str(Path(model_config.model_cache_dir) / "base_model")
    else:
        download_dir = str(model_config.model_cache_dir)

    potential_max_model_length_config_names = [
        "max_position_embeddings",
        "max_sequence_length",
        "model_max_length",
        "sliding_window",
        "sliding_window_size",
        "n_positions",
    ]
    true_max_model_len_candidates: list[int] = list()
    for config_name in potential_max_model_length_config_names:
        if hasattr(hf_model_config, config_name):
            model_len = getattr(hf_model_config, config_name)
            if model_len is not None:
                true_max_model_len_candidates.append(model_len)

    if len(true_max_model_len_candidates) > 0:
        true_max_model_len = min(true_max_model_len_candidates)
    else:
        true_max_model_len = 5_000

    clear_vllm()

    executor_backend = "ray" if torch.cuda.device_count() > 1 else "mp"

    try:
        model = LLM(
            model=model_id,
            tokenizer=model_id,
            gpu_memory_utilization=0.95,
            max_model_len=min(true_max_model_len, 5_000),
            download_dir=download_dir,
            trust_remote_code=benchmark_config.trust_remote_code,
            revision=model_config.revision,
            seed=4242,
            distributed_executor_backend=executor_backend,
            tensor_parallel_size=torch.cuda.device_count(),
            disable_custom_all_reduce=True,
            quantization=quantization,
            dtype=dtype,
            enforce_eager=True,
            max_logprobs=MAX_LOGPROBS if output_scores else None,
            # TEMP: Prefix caching isn't supported with sliding window in vLLM yet,
            # so we disable it for now
            enable_prefix_caching=False,
            enable_lora=model_config.adapter_base_model_id is not None,
            max_lora_rank=256,
        )
    except (ValueError, OSError) as e:
        if "awaiting a review from the repo authors" in str(e):
            raise InvalidModel(
                f"The model {model_id!r} is awaiting a review from the repository "
                "authors. Please try again later."
            )
        elif "trust_remote_code" in str(e):
            raise InvalidModel(
                f"Loading the model {model_id!r} needs to trust remote code. "
                "If you trust the suppliers of this model, then you can enable "
                "this by setting the `--trust-remote-code` flag."
            )
        raise InvalidModel(
            f"The model {model_id!r} could not be loaded. The error was {e!r}."
        )

    model._run_engine = MethodType(_run_engine_with_fixed_progress_bars, model)
    model.config = hf_model_config

    tokenizer = load_tokenizer(
        model_id=model_config.model_id,
        revision=model_config.revision,
        adapter_base_model_id=model_config.adapter_base_model_id,
        trust_remote_code=benchmark_config.trust_remote_code,
        model_max_length=true_max_model_len,
        model_cache_dir=model_config.model_cache_dir,
        token=benchmark_config.api_key or os.getenv("HUGGINGFACE_API_KEY") or True,
    )

    return model, tokenizer


def load_tokenizer(
docs
    model_id: str,
    revision: str,
    adapter_base_model_id: str | None,
    trust_remote_code: bool,
    model_max_length: int,
    model_cache_dir: str,
    token: str | bool,
) -> "PreTrainedTokenizer":
    """Load the tokenizer.

    Args:
        model_id:
            The model identifier.
        revision:
            The revision of the model.
        adapter_base_model_id:
            The base model ID for the adapter model. Can be None if the model is not an
            adapter model.
        trust_remote_code:
            Whether to trust remote code.
        model_max_length:
            The maximum length of the model.
        model_cache_dir:
            The cache directory for the model.
        token:
            The Hugging Face API token.

    Returns:
        The loaded tokenizer.
    """
    config = AutoConfig.from_pretrained(
        adapter_base_model_id or model_id,
        revision=revision,
        cache_dir=model_cache_dir,
        token=token,
        trust_remote_code=trust_remote_code,
    )
    num_retries = 5
    for _ in range(num_retries):
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                use_fast=True,
                verbose=False,
                trust_remote_code=trust_remote_code,
                padding_side="left",
                truncation_side="left",
                model_max_length=model_max_length,
                config=config,
                token=token,
            )
            break
        except (json.JSONDecodeError, OSError, TypeError) as e:
            if adapter_base_model_id is None or model_id == adapter_base_model_id:
                raise InvalidModel(
                    f"Could not load tokenizer for model {model_id!r}. The error was "
                    f"{str(e)}."
                )
            logger.debug(
                f"Could not load tokenizer for {model_id!r}. Falling back to "
                f"{adapter_base_model_id!r}."
            )
            model_id = adapter_base_model_id
        except (TimeoutError, RequestError):
            logger.info(f"Couldn't load tokenizer for {model_id!r}. Retrying.")
            sleep(5)
            continue
    else:
        raise InvalidModel(
            f"Could not load tokenizer for model {model_id!r} after {num_retries} "
            "attempts."
        )

    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    return tokenizer


def _run_engine_with_fixed_progress_bars(
    self: "LLM", use_tqdm: bool
) -> list["RequestOutput"]:
    if use_tqdm:
        num_requests = self.llm_engine.get_num_unfinished_requests()
        pbar = tqdm(
            total=num_requests, leave=False, disable=hasattr(sys, "_called_from_test")
        )
    else:
        pbar = None

    # Run the engine.
    outputs: list["RequestOutput"] = list()
    while self.llm_engine.has_unfinished_requests():
        step_outputs = self.llm_engine.step()
        for output in step_outputs:
            if output.finished:
                outputs.append(output)
                if pbar is not None:
                    pbar.update(1)

    if pbar is not None:
        pbar.close()

    # Sort the outputs by request ID. This is necessary because some requests may be
    # finished earlier than its previous requests.
    outputs = sorted(outputs, key=lambda x: int(x.request_id))

    return outputs


def clear_vllm() -> None:
docs
    """Clear the GPU memory used by the vLLM model, enabling re-initialisation."""
    try:
        destroy_model_parallel()
    except ImportError:
        pass
    clear_memory()
    if ray.is_initialized():
        ray.shutdown()


def get_end_of_reasoning_token_id(
docs
    model: "LLM", tokenizer: "PreTrainedTokenizer"
) -> int | None:
    """Get the end of reasoning token ID for a generative model.

    This assumes that the reasoning token is of the form <X> and that the end of
    reasoning token is </X> (for X being any string without spaces).

    Args:
        model:
            The vLLM model.
        tokenizer:
            The tokenizer.

    Returns:
        The end of reasoning token ID, or None if it could not be found.
    """
    if tokenizer.chat_template is None:
        prompt = "What is your name?"
    else:
        prompt = tokenizer.apply_chat_template(
            conversation=[dict(role="user", content="What is your name?")],
            add_generation_prompt=True,
            tokenize=False,
        )

    # Generate a completion and remove the BOS token from it, to not confuse it with the
    # potential reasoning token
    completion = (
        model.generate(
            prompts=[prompt],
            sampling_params=SamplingParams(max_tokens=3, temperature=0.0),
            use_tqdm=False,
        )[0]
        .outputs[0]
        .text
    )
    if tokenizer.bos_token is not None:
        completion = completion.replace(tokenizer.bos_token, "").strip()

    # If it doesn't contain a reasoning token, we can't find the end of reasoning token
    match = re.search(pattern=r"<\w+>", string=completion)
    if match is None:
        log_once(
            message=(
                "Could not find a reasoning token, so assuming the model is not a "
                "reasoning model."
            ),
            level=logging.DEBUG,
        )
        return None

    # Check that the found reasoning token and its associated end-of-reasoning tokens
    # are both special tokens
    reasoning_token = match.group()
    end_of_reasoning_token = f"</{reasoning_token[1:-1]}>"
    special_tokens = [
        decoder_token.content
        for decoder_token in tokenizer.added_tokens_decoder.values()
    ]
    special_tokens.extend(
        [encoder_token for encoder_token in tokenizer.added_tokens_encoder.keys()]
    )
    special_tokens.extend(tokenizer.all_special_tokens)
    if (
        reasoning_token not in special_tokens
        or end_of_reasoning_token not in special_tokens
    ):
        log_once(
            message=(
                f"Detected reasoning token {reasoning_token!r} and end of reasoning "
                f"token {end_of_reasoning_token!r}, but one of them is not registered "
                "as a special token, so assuming it is not a real reasoning token."
            ),
            level=logging.DEBUG,
        )
        return None

    log_once(
        message=f"Detected reasoning token {reasoning_token!r}.", level=logging.DEBUG
    )

    # Encode the end of reasoning token and return its ID
    end_of_reasoning_token_id = tokenizer.encode(
        text=end_of_reasoning_token, add_special_tokens=False
    )[0]

    return end_of_reasoning_token_id