| from __future__ import annotations |
|
|
| import importlib.metadata |
| import math |
| import os |
| from pathlib import Path |
| from typing import Any, Iterable, Optional, Sequence |
|
|
| from .constants import ( |
| ARCHITECTURE, |
| DEFAULT_DECODER_PAD_TO_MULTIPLE_OF, |
| DEFAULT_DOCUMENT_MAX_LENGTH, |
| DEFAULT_ENCODER_CHUNK_SIZE, |
| DEFAULT_INSTRUCTION, |
| DEFAULT_MAX_MODEL_LEN, |
| DEFAULT_QUERY_MAX_LENGTH, |
| DEFAULT_SYSTEM_INSTRUCTION, |
| MODEL_ID, |
| NO_TOKEN_ID, |
| PLUGIN_NAME, |
| SUPPORTED_ENCODER_CHUNK_SIZES, |
| TEXT_MODALITY, |
| YES_TOKEN_ID, |
| decoder_text, |
| encoder_text, |
| parse_encoder_chunk_size, |
| validate_answer_tokens, |
| ) |
|
|
|
|
| TESTED_VLLM_VERSION = "0.19.1" |
|
|
|
|
| def _positive_int(value: int, name: str) -> int: |
| parsed = int(value) |
| if parsed <= 0: |
| raise ValueError(f"{name} must be a positive integer.") |
| return parsed |
|
|
|
|
| def _sigmoid(value: float) -> float: |
| if value >= 0: |
| scale = math.exp(-value) |
| return 1.0 / (1.0 + scale) |
| scale = math.exp(value) |
| return scale / (1.0 + scale) |
|
|
|
|
| def build_hf_overrides(encoder_chunk_size: int) -> dict[str, object]: |
| return { |
| "architectures": [ARCHITECTURE], |
| "num_labels": 1, |
| "yes_token_id": YES_TOKEN_ID, |
| "no_token_id": NO_TOKEN_ID, |
| "encoder_chunk_size": encoder_chunk_size, |
| "decoder_pad_to_multiple_of": DEFAULT_DECODER_PAD_TO_MULTIPLE_OF, |
| "problem_type": "regression", |
| } |
|
|
|
|
| def _enable_plugin() -> None: |
| allowed = os.environ.get("VLLM_PLUGINS") |
| if allowed is None: |
| os.environ["VLLM_PLUGINS"] = PLUGIN_NAME |
| return |
| names = {item.strip() for item in allowed.split(",") if item.strip()} |
| if PLUGIN_NAME not in names: |
| raise RuntimeError( |
| f"VLLM_PLUGINS={allowed!r} excludes {PLUGIN_NAME!r}. Add the plugin " |
| "name or unset VLLM_PLUGINS." |
| ) |
|
|
|
|
| def check_runtime() -> None: |
| installed_vllm = importlib.metadata.version("vllm") |
| if installed_vllm != TESTED_VLLM_VERSION: |
| raise RuntimeError( |
| f"This adapter requires vLLM {TESTED_VLLM_VERSION}; " |
| f"found {installed_vllm}." |
| ) |
|
|
| discovered = { |
| entry.name: entry.value |
| for entry in importlib.metadata.entry_points(group="vllm.general_plugins") |
| } |
| if PLUGIN_NAME not in discovered: |
| raise RuntimeError( |
| f"vLLM plugin {PLUGIN_NAME!r} is not installed. Install the " |
| "vllm_support package before creating the reranker." |
| ) |
|
|
| from vllm.model_executor.models import ModelRegistry |
| from vllm.plugins import load_general_plugins |
|
|
| load_general_plugins() |
| if ARCHITECTURE not in set(ModelRegistry.get_supported_archs()): |
| raise RuntimeError(f"{ARCHITECTURE} was not registered in ModelRegistry.") |
|
|
|
|
| class KaLMVLLMReranker: |
| """Single-GPU vLLM adapter preserving the original KaLM score contract.""" |
|
|
| def __init__( |
| self, |
| model: str | Path = MODEL_ID, |
| *, |
| query_max_length: int = DEFAULT_QUERY_MAX_LENGTH, |
| document_max_length: int = DEFAULT_DOCUMENT_MAX_LENGTH, |
| encoder_chunk_size: object = DEFAULT_ENCODER_CHUNK_SIZE, |
| dtype: str = "bfloat16", |
| tensor_parallel_size: int = 1, |
| max_model_len: int = DEFAULT_MAX_MODEL_LEN, |
| gpu_memory_utilization: float = 0.85, |
| batch_size: int = 32, |
| instruction: str = DEFAULT_INSTRUCTION, |
| system_instruction: str = DEFAULT_SYSTEM_INSTRUCTION, |
| skip_runtime_check: bool = False, |
| ) -> None: |
| self.model = str(model) |
| self.query_max_length = _positive_int( |
| query_max_length, "query_max_length" |
| ) |
| self.document_max_length = _positive_int( |
| document_max_length, "document_max_length" |
| ) |
| self.encoder_chunk_size = parse_encoder_chunk_size(encoder_chunk_size) |
| self.dtype = str(dtype) |
| self.tensor_parallel_size = _positive_int( |
| tensor_parallel_size, "tensor_parallel_size" |
| ) |
| if self.tensor_parallel_size != 1: |
| raise ValueError( |
| "The published adapter supports tensor_parallel_size=1 only." |
| ) |
| self.max_model_len = _positive_int(max_model_len, "max_model_len") |
| self.gpu_memory_utilization = float(gpu_memory_utilization) |
| if not 0 < self.gpu_memory_utilization <= 1: |
| raise ValueError("gpu_memory_utilization must be in the interval (0, 1].") |
| self.batch_size = _positive_int(batch_size, "batch_size") |
| if not isinstance(instruction, str) or not isinstance( |
| system_instruction, str |
| ): |
| raise TypeError("instruction and system_instruction must be strings.") |
| self.instruction = instruction |
| self.system_instruction = system_instruction |
|
|
| _enable_plugin() |
| if not skip_runtime_check: |
| check_runtime() |
|
|
| from transformers import AutoTokenizer |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained( |
| self.model, |
| trust_remote_code=True, |
| ) |
| validate_answer_tokens(self.tokenizer) |
|
|
| from vllm import LLM |
| from vllm.pooling_params import PoolingParams |
|
|
| self.pooling_params = PoolingParams(use_activation=False) |
| self.llm = LLM( |
| model=self.model, |
| runner="pooling", |
| trust_remote_code=True, |
| hf_overrides=build_hf_overrides(self.encoder_chunk_size), |
| dtype=self.dtype, |
| tensor_parallel_size=1, |
| max_model_len=self.max_model_len, |
| gpu_memory_utilization=self.gpu_memory_utilization, |
| enforce_eager=True, |
| limit_mm_per_prompt={TEXT_MODALITY: 1}, |
| ) |
|
|
| def close(self) -> None: |
| llm = getattr(self, "llm", None) |
| if llm is None: |
| return |
| engine = getattr(llm, "llm_engine", None) |
| engine_core = getattr(engine, "engine_core", None) |
| shutdown = getattr(engine_core, "shutdown", None) |
| if callable(shutdown): |
| shutdown() |
| self.llm = None |
|
|
| def __enter__(self) -> "KaLMVLLMReranker": |
| return self |
|
|
| def __exit__(self, exc_type, exc_value, traceback) -> None: |
| self.close() |
|
|
| def _encoder_ids(self, document: str) -> tuple[str, list[int]]: |
| text = encoder_text(document) |
| token_ids = self.tokenizer( |
| text, |
| add_special_tokens=False, |
| truncation=True, |
| max_length=self.document_max_length, |
| )["input_ids"] |
| if not token_ids: |
| raise ValueError("Encoded document prompt is empty.") |
| return text, list(token_ids) |
|
|
| def _decoder_ids(self, query: str, instruction: str) -> tuple[str, list[int]]: |
| text = decoder_text( |
| self.tokenizer, |
| query, |
| instruction=instruction, |
| system_instruction=self.system_instruction, |
| query_max_length=self.query_max_length, |
| ) |
| token_ids = self.tokenizer.encode(text, add_special_tokens=False) |
| if not token_ids: |
| raise ValueError("Encoded decoder prompt is empty.") |
| return text, list(token_ids) |
|
|
| def _prompt(self, query: str, document: str, instruction: str): |
| from vllm.inputs import ExplicitEncoderDecoderPrompt, TokensPrompt |
|
|
| encoder_prompt, encoder_ids = self._encoder_ids(document) |
| decoder_prompt, decoder_ids = self._decoder_ids(query, instruction) |
| return ExplicitEncoderDecoderPrompt( |
| encoder_prompt=TokensPrompt( |
| prompt_token_ids=encoder_ids, |
| prompt=encoder_prompt, |
| multi_modal_data={TEXT_MODALITY: [encoder_prompt]}, |
| ), |
| decoder_prompt=TokensPrompt( |
| prompt_token_ids=decoder_ids, |
| prompt=decoder_prompt, |
| ), |
| ) |
|
|
| @staticmethod |
| def _validate_pairs( |
| pairs: Sequence[tuple[str, str]], |
| ) -> list[tuple[str, str]]: |
| if isinstance(pairs, (str, bytes)) or not isinstance(pairs, Sequence): |
| raise TypeError("pairs must be a sequence of (query, document) pairs.") |
| validated: list[tuple[str, str]] = [] |
| for index, pair in enumerate(pairs): |
| if ( |
| isinstance(pair, (str, bytes)) |
| or not isinstance(pair, Sequence) |
| or len(pair) != 2 |
| ): |
| raise ValueError(f"pairs[{index}] must contain exactly two strings.") |
| query, document = pair |
| if not isinstance(query, str) or not isinstance(document, str): |
| raise TypeError(f"pairs[{index}] must contain exactly two strings.") |
| validated.append((query, document)) |
| return validated |
|
|
| @staticmethod |
| def _margins_from_outputs(outputs: Iterable[Any]) -> list[float]: |
| margins: list[float] = [] |
| for output in outputs: |
| values = output.outputs.probs |
| if len(values) != 1: |
| raise RuntimeError(f"Expected one raw margin, got {values}.") |
| margin = float(values[0]) |
| if not math.isfinite(margin): |
| raise RuntimeError(f"vLLM returned a non-finite margin: {margin}.") |
| margins.append(margin) |
| return margins |
|
|
| def predict( |
| self, |
| pairs: Sequence[tuple[str, str]], |
| *, |
| instruction: Optional[str] = None, |
| return_margin: bool = False, |
| ) -> list[float] | list[dict[str, float]]: |
| validated_pairs = self._validate_pairs(pairs) |
| if not validated_pairs: |
| return [] |
| effective_instruction = self.instruction if instruction is None else instruction |
| if not isinstance(effective_instruction, str): |
| raise TypeError("instruction must be a string or None.") |
|
|
| margins: list[float] = [] |
| for start in range(0, len(validated_pairs), self.batch_size): |
| batch = validated_pairs[start : start + self.batch_size] |
| prompts = [ |
| self._prompt(query, document, effective_instruction) |
| for query, document in batch |
| ] |
| if self.llm is None: |
| raise RuntimeError("The reranker has been closed.") |
| outputs = self.llm.classify( |
| prompts, |
| pooling_params=self.pooling_params, |
| use_tqdm=False, |
| ) |
| margins.extend(self._margins_from_outputs(outputs)) |
|
|
| scores = [_sigmoid(margin) for margin in margins] |
| if return_margin: |
| return [ |
| {"score": score, "margin": margin} |
| for score, margin in zip(scores, margins) |
| ] |
| return scores |
|
|
| def rank( |
| self, |
| query: str, |
| documents: Sequence[str], |
| *, |
| instruction: Optional[str] = None, |
| top_k: Optional[int] = None, |
| return_margin: bool = False, |
| ) -> list[dict[str, float | int]]: |
| if not isinstance(query, str): |
| raise TypeError("query must be a string.") |
| if isinstance(documents, (str, bytes)) or not isinstance(documents, Sequence): |
| raise TypeError("documents must be a sequence of strings.") |
| if any(not isinstance(document, str) for document in documents): |
| raise TypeError("every document must be a string.") |
| if top_k is not None: |
| top_k = int(top_k) |
| if top_k < 0: |
| raise ValueError("top_k must be non-negative or None.") |
|
|
| predictions = self.predict( |
| [(query, document) for document in documents], |
| instruction=instruction, |
| return_margin=return_margin, |
| ) |
| rankings: list[dict[str, float | int]] = [] |
| for corpus_id, prediction in enumerate(predictions): |
| if return_margin: |
| assert isinstance(prediction, dict) |
| item: dict[str, float | int] = { |
| "corpus_id": corpus_id, |
| "score": prediction["score"], |
| "margin": prediction["margin"], |
| } |
| else: |
| assert isinstance(prediction, float) |
| item = {"corpus_id": corpus_id, "score": prediction} |
| rankings.append(item) |
| rankings.sort(key=lambda item: float(item["score"]), reverse=True) |
| return rankings if top_k is None else rankings[:top_k] |
|
|
|
|
| KaLMVLLMOfflineReranker = KaLMVLLMReranker |
|
|
| __all__ = [ |
| "KaLMVLLMOfflineReranker", |
| "KaLMVLLMReranker", |
| "SUPPORTED_ENCODER_CHUNK_SIZES", |
| "build_hf_overrides", |
| "check_runtime", |
| "parse_encoder_chunk_size", |
| ] |
|
|
