Update vLLM support and remove build artifacts

vllm_support/README.md

@@ -1,7 +1,7 @@
-# KaLM-Reranker-V1-Small vLLM Support
+# KaLM-Reranker-V1-Nano vLLM Support
 
 This directory contains the experimental vLLM 0.19.1 adapter for
-`KaLM-Embedding/KaLM-Reranker-V1-Small`. It supports offline Python and CLI
+`KaLM-Embedding/KaLM-Reranker-V1-Nano`. It supports offline Python and CLI
 reranking plus an optional FastAPI service.
 
 The adapter does not modify or retrain the checkpoint. It reads the original
@@ -35,9 +35,9 @@ conda activate kalm-vllm
 pip install "vllm==0.19.1" "transformers==5.6.2"
 pip install "fastapi>=0.136,<0.137" "uvicorn>=0.46,<0.47"
 
-hf download KaLM-Embedding/KaLM-Reranker-V1-Small \
-  --local-dir ./KaLM-Reranker-V1-Small
-pip install ./KaLM-Reranker-V1-Small/vllm_support --no-deps
+hf download KaLM-Embedding/KaLM-Reranker-V1-Nano \
+  --local-dir ./KaLM-Reranker-V1-Nano
+pip install ./KaLM-Reranker-V1-Nano/vllm_support --no-deps
 export VLLM_PLUGINS=kalm_t5gemma2
 ```
 
@@ -45,10 +45,10 @@ The model can also be loaded directly by its Hugging Face ID. In that case,
 only download the `vllm_support` directory before installing the plugin:
 
 ```bash
-hf download KaLM-Embedding/KaLM-Reranker-V1-Small \
+hf download KaLM-Embedding/KaLM-Reranker-V1-Nano \
   --include "vllm_support/**" \
-  --local-dir ./KaLM-Reranker-V1-Small
-pip install ./KaLM-Reranker-V1-Small/vllm_support --no-deps
+  --local-dir ./KaLM-Reranker-V1-Nano
+pip install ./KaLM-Reranker-V1-Nano/vllm_support --no-deps
 ```
 
 ## Offline Python API
@@ -64,7 +64,7 @@ documents = [
 
 pairs = [(query, document) for document in documents]
 with KaLMVLLMReranker(
-    "KaLM-Embedding/KaLM-Reranker-V1-Small",
+    "KaLM-Embedding/KaLM-Reranker-V1-Nano",
     query_max_length=512,
     document_max_length=1024,
     encoder_chunk_size=4,
@@ -97,7 +97,7 @@ Score JSONL input:
 
 ```bash
 kalm-vllm-rerank \
-  --input-jsonl ./KaLM-Reranker-V1-Small/vllm_support/examples/sample_pairs.jsonl \
+  --input-jsonl ./KaLM-Reranker-V1-Nano/vllm_support/examples/sample_pairs.jsonl \
   --output-jsonl ./scores.jsonl \
   --return-margin
 ```
@@ -114,7 +114,7 @@ Start one model instance:
 kalm-vllm-serve \
   --host 0.0.0.0 \
   --port 8000 \
-  --model KaLM-Embedding/KaLM-Reranker-V1-Small \
+  --model KaLM-Embedding/KaLM-Reranker-V1-Nano \
   --encoder-chunk-size 4
 ```
 
@@ -123,7 +123,7 @@ variables:
 
 ```bash
 CUDA_VISIBLE_DEVICES=0 PORT=8000 \
-  ./KaLM-Reranker-V1-Small/vllm_support/examples/start_online_server.sh
+  ./KaLM-Reranker-V1-Nano/vllm_support/examples/start_online_server.sh
 ```
 
 In a second terminal, check health and send built-in demo requests:
@@ -140,7 +140,7 @@ one query against multiple documents:
 ```bash
 kalm-vllm-client \
   --endpoint rerank \
-  --json-file ./KaLM-Reranker-V1-Small/vllm_support/examples/rerank_request.json \
+  --json-file ./KaLM-Reranker-V1-Nano/vllm_support/examples/rerank_request.json \
   --return-margin \
   --top-k 10
 ```
@@ -150,7 +150,7 @@ Use `/score` for a batch of independent query-document pairs:
 ```bash
 kalm-vllm-client \
   --endpoint score \
-  --json-file ./KaLM-Reranker-V1-Small/vllm_support/examples/score_request.json \
+  --json-file ./KaLM-Reranker-V1-Nano/vllm_support/examples/score_request.json \
   --return-margin
 ```
 
@@ -247,7 +247,7 @@ Reinstall the package and ensure the environment variable includes its entry
 point name:
 
 ```bash
-pip install ./KaLM-Reranker-V1-Small/vllm_support --no-deps --force-reinstall
+pip install ./KaLM-Reranker-V1-Nano/vllm_support --no-deps --force-reinstall
 export VLLM_PLUGINS=kalm_t5gemma2
 ```
 
@@ -258,7 +258,7 @@ sensitive.
 
 **The tokenizer check fails**
 
-Confirm that the tokenizer belongs to this Small checkpoint. The adapter
+Confirm that the tokenizer belongs to this Nano checkpoint. The adapter
 requires `yes -> 4443` and `no -> 1904`.
 
 **CUDA runs out of memory**

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/__init__.py

@@ -1,23 +0,0 @@
-from __future__ import annotations
-
-from .constants import (
-    ARCHITECTURE,
-    MODEL_ID,
-    PLUGIN_NAME,
-    SUPPORTED_ENCODER_CHUNK_SIZES,
-    TEXT_MODALITY,
-)
-from .reranker import KaLMVLLMOfflineReranker, KaLMVLLMReranker
-
-__version__ = "0.1.0"
-
-__all__ = [
-    "ARCHITECTURE",
-    "KaLMVLLMOfflineReranker",
-    "KaLMVLLMReranker",
-    "MODEL_ID",
-    "PLUGIN_NAME",
-    "SUPPORTED_ENCODER_CHUNK_SIZES",
-    "TEXT_MODALITY",
-    "__version__",
-]

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/cli.py

@@ -1,172 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import json
-import sys
-from collections import OrderedDict
-from pathlib import Path
-from typing import Any, Iterable, TextIO
-
-from .constants import (
-    DEFAULT_DOCUMENT_MAX_LENGTH,
-    DEFAULT_ENCODER_CHUNK_SIZE,
-    DEFAULT_MAX_MODEL_LEN,
-    DEFAULT_QUERY_MAX_LENGTH,
-    MODEL_ID,
-    SAMPLE_DOCUMENTS,
-    SAMPLE_QUERY,
-    SUPPORTED_ENCODER_CHUNK_SIZES,
-    parse_encoder_chunk_size,
-)
-from .reranker import KaLMVLLMReranker
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        description="Offline KaLM-Reranker-V1-Nano scoring with vLLM 0.19.1.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument("--input-jsonl", type=Path)
-    parser.add_argument("--output-jsonl", type=Path)
-    parser.add_argument("--model", default=MODEL_ID)
-    parser.add_argument("--query-max-length", type=int, default=DEFAULT_QUERY_MAX_LENGTH)
-    parser.add_argument(
-        "--document-max-length", type=int, default=DEFAULT_DOCUMENT_MAX_LENGTH
-    )
-    parser.add_argument(
-        "--encoder-chunk-size",
-        default=str(DEFAULT_ENCODER_CHUNK_SIZE),
-        help=f"One of {sorted(SUPPORTED_ENCODER_CHUNK_SIZES)}.",
-    )
-    parser.add_argument("--max-model-len", type=int, default=DEFAULT_MAX_MODEL_LEN)
-    parser.add_argument("--batch-size", type=int, default=32)
-    parser.add_argument("--return-margin", action="store_true")
-    parser.add_argument("--top-k", type=int)
-    parser.add_argument("--dtype", default="bfloat16")
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85)
-    parser.add_argument("--tensor-parallel-size", type=int, default=1)
-    return parser
-
-
-def _read_jsonl(path: Path) -> list[dict[str, Any]]:
-    rows: list[dict[str, Any]] = []
-    with path.open("r", encoding="utf-8") as handle:
-        for line_no, line in enumerate(handle, start=1):
-            if not line.strip():
-                continue
-            try:
-                row = json.loads(line)
-            except json.JSONDecodeError as error:
-                raise ValueError(f"{path}:{line_no}: invalid JSON.") from error
-            if not isinstance(row, dict):
-                raise ValueError(f"{path}:{line_no}: expected a JSON object.")
-            if not isinstance(row.get("query"), str):
-                raise ValueError(f"{path}:{line_no}: 'query' must be a string.")
-            if not isinstance(row.get("document"), str):
-                raise ValueError(f"{path}:{line_no}: 'document' must be a string.")
-            if "instruction" in row and not isinstance(row["instruction"], str):
-                raise ValueError(f"{path}:{line_no}: 'instruction' must be a string.")
-            rows.append(row)
-    return rows
-
-
-def _write_jsonl(rows: Iterable[dict[str, Any]], output: TextIO) -> None:
-    for row in rows:
-        output.write(json.dumps(row, ensure_ascii=False) + "\n")
-    output.flush()
-
-
-def _score_rows(
-    reranker: KaLMVLLMReranker,
-    rows: list[dict[str, Any]],
-    *,
-    return_margin: bool,
-    top_k: int | None,
-) -> list[dict[str, Any]]:
-    grouped: OrderedDict[str | None, list[tuple[int, dict[str, Any]]]] = OrderedDict()
-    for index, row in enumerate(rows):
-        grouped.setdefault(row.get("instruction"), []).append((index, row))
-
-    scored: dict[int, dict[str, Any]] = {}
-    for instruction, items in grouped.items():
-        predictions = reranker.predict(
-            [(row["query"], row["document"]) for _, row in items],
-            instruction=instruction,
-            return_margin=True,
-        )
-        for (index, row), prediction in zip(items, predictions):
-            assert isinstance(prediction, dict)
-            result = {
-                "id": row.get("id", index),
-                "query": row["query"],
-                "document": row["document"],
-                "score": prediction["score"],
-            }
-            if "instruction" in row:
-                result["instruction"] = row["instruction"]
-            if return_margin:
-                result["margin"] = prediction["margin"]
-            scored[index] = result
-
-    ordered = [scored[index] for index in range(len(rows))]
-    if top_k is None:
-        return ordered
-    if top_k < 0:
-        raise ValueError("--top-k must be non-negative.")
-    by_query: OrderedDict[str, list[dict[str, Any]]] = OrderedDict()
-    for row in ordered:
-        by_query.setdefault(str(row["query"]), []).append(row)
-    output: list[dict[str, Any]] = []
-    for group in by_query.values():
-        group.sort(key=lambda item: float(item["score"]), reverse=True)
-        output.extend(group[:top_k])
-    return output
-
-
-def main() -> int:
-    args = build_parser().parse_args()
-    with KaLMVLLMReranker(
-        args.model,
-        query_max_length=args.query_max_length,
-        document_max_length=args.document_max_length,
-        encoder_chunk_size=parse_encoder_chunk_size(args.encoder_chunk_size),
-        max_model_len=args.max_model_len,
-        batch_size=args.batch_size,
-        dtype=args.dtype,
-        gpu_memory_utilization=args.gpu_memory_utilization,
-        tensor_parallel_size=args.tensor_parallel_size,
-    ) as reranker:
-        if args.input_jsonl is None:
-            rankings = reranker.rank(
-                SAMPLE_QUERY,
-                SAMPLE_DOCUMENTS,
-                top_k=args.top_k,
-                return_margin=args.return_margin,
-            )
-            rows = [
-                {
-                    "query": SAMPLE_QUERY,
-                    "document": SAMPLE_DOCUMENTS[int(item["corpus_id"])],
-                    **item,
-                }
-                for item in rankings
-            ]
-        else:
-            rows = _score_rows(
-                reranker,
-                _read_jsonl(args.input_jsonl),
-                return_margin=args.return_margin,
-                top_k=args.top_k,
-            )
-
-    if args.output_jsonl is None:
-        _write_jsonl(rows, sys.stdout)
-    else:
-        args.output_jsonl.parent.mkdir(parents=True, exist_ok=True)
-        with args.output_jsonl.open("w", encoding="utf-8") as handle:
-            _write_jsonl(rows, handle)
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/client.py

@@ -1,125 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import json
-import sys
-import urllib.error
-import urllib.request
-from pathlib import Path
-from typing import Any
-
-from .constants import SAMPLE_DOCUMENTS, SAMPLE_QUERY
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        description="Call the KaLM vLLM FastAPI service.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument("--base-url", default="http://127.0.0.1:8000")
-    parser.add_argument("--health", action="store_true")
-    parser.add_argument("--endpoint", choices=("rerank", "score"), default="rerank")
-    parser.add_argument("--json-file", type=Path)
-    parser.add_argument("--return-margin", action="store_true")
-    parser.add_argument("--top-k", type=int)
-    parser.add_argument("--timeout", type=float, default=600.0)
-    return parser
-
-
-def request_json(
-    method: str,
-    url: str,
-    *,
-    payload: dict[str, Any] | None = None,
-    timeout: float = 600.0,
-) -> dict[str, Any]:
-    data = None
-    headers = {"Accept": "application/json"}
-    if payload is not None:
-        data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
-        headers["Content-Type"] = "application/json"
-    request = urllib.request.Request(url, data=data, headers=headers, method=method)
-    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
-    try:
-        with opener.open(request, timeout=timeout) as response:
-            raw = response.read().decode("utf-8")
-    except urllib.error.HTTPError as error:
-        detail = error.read().decode("utf-8", errors="replace")
-        raise RuntimeError(f"HTTP {error.code} from {url}: {detail}") from error
-    return json.loads(raw)
-
-
-def _demo_payload(
-    endpoint: str,
-    return_margin: bool,
-    top_k: int | None,
-) -> dict[str, Any]:
-    if endpoint == "rerank":
-        return {
-            "query": SAMPLE_QUERY,
-            "documents": list(SAMPLE_DOCUMENTS),
-            "top_k": top_k,
-            "return_margin": return_margin,
-        }
-    return {
-        "pairs": [
-            {
-                "id": "positive",
-                "query": SAMPLE_QUERY,
-                "document": SAMPLE_DOCUMENTS[0],
-            },
-            {
-                "id": "negative",
-                "query": SAMPLE_QUERY,
-                "document": SAMPLE_DOCUMENTS[1],
-            },
-        ],
-        "return_margin": return_margin,
-    }
-
-
-def _load_payload(path: Path) -> dict[str, Any]:
-    with path.open("r", encoding="utf-8") as handle:
-        payload = json.load(handle)
-    if not isinstance(payload, dict):
-        raise ValueError(f"{path} must contain one JSON object.")
-    return payload
-
-
-def _request_payload(args: argparse.Namespace) -> dict[str, Any]:
-    payload = (
-        _load_payload(args.json_file)
-        if args.json_file is not None
-        else _demo_payload(args.endpoint, args.return_margin, args.top_k)
-    )
-    if args.return_margin:
-        payload["return_margin"] = True
-    if args.top_k is not None:
-        if args.endpoint != "rerank":
-            raise ValueError("--top-k is only valid with --endpoint rerank.")
-        if args.top_k < 0:
-            raise ValueError("--top-k must be non-negative.")
-        payload["top_k"] = args.top_k
-    return payload
-
-
-def main() -> int:
-    args = build_parser().parse_args()
-    base_url = args.base_url.rstrip("/")
-    if args.health:
-        response = request_json("GET", f"{base_url}/health", timeout=args.timeout)
-    else:
-        payload = _request_payload(args)
-        response = request_json(
-            "POST",
-            f"{base_url}/{args.endpoint}",
-            payload=payload,
-            timeout=args.timeout,
-        )
-    json.dump(response, sys.stdout, ensure_ascii=False, indent=2)
-    sys.stdout.write("\n")
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/constants.py

@@ -1,105 +0,0 @@
-from __future__ import annotations
-
-from typing import Iterable
-
-
-MODEL_ID = "KaLM-Embedding/KaLM-Reranker-V1-Nano"
-PLUGIN_NAME = "kalm_t5gemma2"
-ARCHITECTURE = "T5Gemma2VllmScoreClassification"
-TEXT_MODALITY = "text"
-
-DEFAULT_INSTRUCTION = "Given a query, retrieve documents that answer the query."
-DEFAULT_SYSTEM_INSTRUCTION = (
-    "Judge whether the Document meets the requirements based on the Query and "
-    'the Instruct provided. Note that the answer can only be "yes" or "no".'
-)
-
-DEFAULT_QUERY_MAX_LENGTH = 512
-DEFAULT_DOCUMENT_MAX_LENGTH = 1024
-DEFAULT_MAX_MODEL_LEN = 2048
-DEFAULT_ENCODER_CHUNK_SIZE = 4
-DEFAULT_DECODER_PAD_TO_MULTIPLE_OF = 8
-SUPPORTED_ENCODER_CHUNK_SIZES = frozenset({1, 2, 4, 8, 16, 32})
-
-YES_TOKEN_TEXT = "yes"
-NO_TOKEN_TEXT = "no"
-YES_TOKEN_ID = 4443
-NO_TOKEN_ID = 1904
-
-SAMPLE_QUERY = "What is the capital of China?"
-SAMPLE_DOCUMENTS = (
-    "The capital of China is Beijing.",
-    "Gravity attracts bodies toward one another.",
-)
-
-
-def parse_encoder_chunk_size(value: object) -> int:
-    if value is None:
-        return DEFAULT_ENCODER_CHUNK_SIZE
-    if isinstance(value, str):
-        normalized = value.strip().lower()
-        if normalized in {"", "none", "null"}:
-            return DEFAULT_ENCODER_CHUNK_SIZE
-        try:
-            parsed = int(normalized)
-        except ValueError as error:
-            raise ValueError(_chunk_size_error(value)) from error
-    else:
-        parsed = int(value)
-    if parsed not in SUPPORTED_ENCODER_CHUNK_SIZES:
-        raise ValueError(_chunk_size_error(value))
-    return parsed
-
-
-def _chunk_size_error(value: object) -> str:
-    return (
-        "encoder_chunk_size must be one of "
-        f"{sorted(SUPPORTED_ENCODER_CHUNK_SIZES)}; got {value!r}."
-    )
-
-
-def encoder_text(document: str) -> str:
-    return f"<Document>: {document}"
-
-
-def decoder_text(
-    tokenizer,
-    query: str,
-    *,
-    instruction: str = DEFAULT_INSTRUCTION,
-    system_instruction: str = DEFAULT_SYSTEM_INSTRUCTION,
-    query_max_length: int = DEFAULT_QUERY_MAX_LENGTH,
-) -> str:
-    query_ids = tokenizer(
-        query,
-        add_special_tokens=False,
-        truncation=True,
-        max_length=query_max_length,
-    )["input_ids"]
-    truncated_query = tokenizer.decode(
-        query_ids,
-        skip_special_tokens=False,
-        clean_up_tokenization_spaces=False,
-    )
-    return (
-        "<bos><start_of_turn>user\n"
-        f"{system_instruction}\n\n"
-        f"<Instruct>: {instruction}\n"
-        f"<Query>: {truncated_query}<end_of_turn>\n"
-        "<start_of_turn>model\n\n\n\n"
-    )
-
-
-def validate_answer_tokens(tokenizer) -> None:
-    expected: Iterable[tuple[str, int]] = (
-        (YES_TOKEN_TEXT, YES_TOKEN_ID),
-        (NO_TOKEN_TEXT, NO_TOKEN_ID),
-    )
-    for text, token_id in expected:
-        actual = tokenizer.encode(text, add_special_tokens=False)
-        if actual != [token_id]:
-            raise RuntimeError(
-                f"Unexpected tokenization for {text!r}: expected {[token_id]}, "
-                f"got {actual}."
-            )
-

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/modeling.py

@@ -1,265 +0,0 @@
-from __future__ import annotations
-
-import math
-import os
-from typing import Iterable
-
-import torch
-import torch.nn as nn
-from transformers.modeling_outputs import BaseModelOutput
-from vllm.model_executor.layers.pooler import DispatchPooler
-from vllm.multimodal import MULTIMODAL_REGISTRY
-
-from .constants import (
-    DEFAULT_DECODER_PAD_TO_MULTIPLE_OF,
-    DEFAULT_ENCODER_CHUNK_SIZE,
-    NO_TOKEN_ID,
-    TEXT_MODALITY,
-    YES_TOKEN_ID,
-)
-from .modeling_score import T5Gemma2ForScoreClassification
-from .processing import (
-    TextEncoderDummyInputsBuilder,
-    TextEncoderProcessingInfo,
-    TextEncoderProcessor,
-)
-
-
-def _as_token_rows(value: object) -> list[torch.Tensor]:
-    if isinstance(value, torch.Tensor):
-        if value.ndim == 1:
-            return [value]
-        if value.ndim == 2:
-            return [row for row in value]
-        raise ValueError(f"encoder_input_ids must be 1D/2D, got {value.shape}.")
-    if isinstance(value, list):
-        return [
-            item.flatten()
-            if isinstance(item, torch.Tensor)
-            else torch.tensor(item, dtype=torch.long)
-            for item in value
-        ]
-    raise TypeError(f"Unsupported encoder_input_ids type: {type(value)!r}")
-
-
-def _split_by_position_zero(
-    input_ids: torch.Tensor,
-    positions: torch.Tensor,
-) -> tuple[list[torch.Tensor], list[int]]:
-    flat_ids = input_ids.flatten()
-    starts = (positions.flatten() == 0).nonzero(as_tuple=False).flatten().tolist()
-    if not starts:
-        return [], []
-    starts.append(int(flat_ids.numel()))
-    rows: list[torch.Tensor] = []
-    last_indices: list[int] = []
-    for start, end in zip(starts[:-1], starts[1:]):
-        if end > start:
-            rows.append(flat_ids[start:end])
-            last_indices.append(end - 1)
-    return rows, last_indices
-
-
-def _debug(message: str) -> None:
-    if os.environ.get("KALM_VLLM_DEBUG") == "1":
-        print(f"[kalm-vllm-debug] {message}", flush=True)
-
-
-@MULTIMODAL_REGISTRY.register_processor(
-    TextEncoderProcessor,
-    info=TextEncoderProcessingInfo,
-    dummy_inputs=TextEncoderDummyInputsBuilder,
-)
-class T5Gemma2VllmScoreClassification(nn.Module):
-    is_pooling_model = True
-    supports_multimodal = True
-    score_type = "cross-encoder"
-    attn_type = "encoder_decoder"
-    default_seq_pooling_type = "LAST"
-    default_tok_pooling_type = "ALL"
-
-    def __init__(self, *, vllm_config, prefix: str = "") -> None:
-        super().__init__()
-        self.vllm_config = vllm_config
-        self.model_config = vllm_config.model_config
-        self.config = self.model_config.hf_config
-        self.config.num_labels = 1
-        self.config.yes_token_id = int(
-            getattr(self.config, "yes_token_id", YES_TOKEN_ID)
-        )
-        self.config.no_token_id = int(
-            getattr(self.config, "no_token_id", NO_TOKEN_ID)
-        )
-        self.config.encoder_chunk_size = getattr(
-            self.config,
-            "encoder_chunk_size",
-            DEFAULT_ENCODER_CHUNK_SIZE,
-        )
-        self.encoder_chunk_size = self.config.encoder_chunk_size
-        self.config.decoder_pad_to_multiple_of = int(
-            getattr(
-                self.config,
-                "decoder_pad_to_multiple_of",
-                DEFAULT_DECODER_PAD_TO_MULTIPLE_OF,
-            )
-        )
-        self.decoder_pad_to_multiple_of = self.config.decoder_pad_to_multiple_of
-        self.pad_token_id = int(getattr(self.config, "pad_token_id", 0) or 0)
-
-        self.score_model = T5Gemma2ForScoreClassification.from_pretrained(
-            self.model_config.model,
-            trust_remote_code=True,
-            dtype=self.model_config.dtype,
-        )
-        self.score_model.config.yes_token_id = self.config.yes_token_id
-        self.score_model.config.no_token_id = self.config.no_token_id
-        self.score_model.config.num_labels = 1
-        self.score_model.config.encoder_chunk_size = self.encoder_chunk_size
-        self.score_model.yes_token_id = self.config.yes_token_id
-        self.score_model.no_token_id = self.config.no_token_id
-        self.score_model.encoder_chunk_size = self.encoder_chunk_size
-        self.score_model._validate_score_config()
-        self.score_model.eval()
-
-        pooler_config = self.model_config.pooler_config
-        assert pooler_config is not None
-        self.pooler = DispatchPooler.for_seq_cls(pooler_config)
-
-    def get_language_model(self):
-        return self
-
-    def get_num_mm_encoder_tokens(self, num_tokens: int) -> int:
-        if self.encoder_chunk_size is None:
-            return int(num_tokens)
-        return int(math.ceil(num_tokens / int(self.encoder_chunk_size)))
-
-    def embed_input_ids(self, input_ids: torch.Tensor, *args, **kwargs) -> torch.Tensor:
-        return self.score_model.get_input_embeddings()(input_ids)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # The semantic wrapper loads the same checkpoint directly.
-        for _ in weights:
-            pass
-        return set(self.state_dict().keys())
-
-    def embed_multimodal(self, **kwargs: object) -> list[torch.Tensor]:
-        if "encoder_input_ids" not in kwargs:
-            raise ValueError(f"Missing {TEXT_MODALITY} encoder_input_ids.")
-        rows = _as_token_rows(kwargs["encoder_input_ids"])
-        device = next(self.score_model.parameters()).device
-        outputs: list[torch.Tensor] = []
-        with torch.inference_mode():
-            for row in rows:
-                input_ids = row.to(device=device, dtype=torch.long).unsqueeze(0)
-                attention_mask = torch.ones_like(input_ids)
-                raw_encoder_outputs = self.score_model.get_encoder()(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    return_dict=True,
-                )
-                hidden = raw_encoder_outputs.last_hidden_state
-                if self.encoder_chunk_size is not None:
-                    hidden, _ = self.score_model._pool_encoder_chunks(
-                        hidden,
-                        attention_mask,
-                        int(self.encoder_chunk_size),
-                    )
-                item = hidden.squeeze(0).contiguous()
-                _debug(f"encoder output shape={tuple(item.shape)}")
-                outputs.append(item)
-        return outputs
-
-    def _pad_decoder_rows(
-        self,
-        rows: list[torch.Tensor],
-        device: torch.device,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        max_len = max(int(row.numel()) for row in rows)
-        if self.decoder_pad_to_multiple_of > 1:
-            multiple = self.decoder_pad_to_multiple_of
-            max_len = int(math.ceil(max_len / multiple) * multiple)
-        batch = torch.full(
-            (len(rows), max_len),
-            self.pad_token_id,
-            dtype=torch.long,
-            device=device,
-        )
-        mask = torch.zeros_like(batch)
-        for index, row in enumerate(rows):
-            row = row.to(device=device, dtype=torch.long)
-            length = int(row.numel())
-            batch[index, :length] = row
-            mask[index, :length] = 1
-        return batch, mask
-
-    @staticmethod
-    def _pad_encoder_outputs(
-        encoder_outputs: list[torch.Tensor],
-        device: torch.device,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        max_len = max(int(item.shape[0]) for item in encoder_outputs)
-        hidden_size = int(encoder_outputs[0].shape[-1])
-        batch = torch.zeros(
-            (len(encoder_outputs), max_len, hidden_size),
-            dtype=encoder_outputs[0].dtype,
-            device=device,
-        )
-        mask = torch.zeros(
-            (len(encoder_outputs), max_len),
-            dtype=torch.long,
-            device=device,
-        )
-        for index, item in enumerate(encoder_outputs):
-            item = item.to(device=device)
-            length = int(item.shape[0])
-            batch[index, :length] = item
-            mask[index, :length] = 1
-        return batch, mask
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        intermediate_tensors=None,
-        inputs_embeds: torch.Tensor | None = None,
-        encoder_outputs: list[torch.Tensor] | torch.Tensor | None = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        if input_ids is None:
-            raise ValueError("Decoder input_ids are required.")
-        decoder_rows, last_indices = _split_by_position_zero(input_ids, positions)
-        hidden = input_ids.new_zeros((input_ids.numel(), 1), dtype=torch.float32)
-        if not decoder_rows:
-            return hidden
-        if encoder_outputs is None:
-            return hidden
-
-        encoder_list = (
-            [item for item in encoder_outputs]
-            if isinstance(encoder_outputs, torch.Tensor)
-            else list(encoder_outputs)
-        )
-        if len(encoder_list) != len(decoder_rows):
-            raise ValueError(
-                "Mismatched encoder/decoder batch sizes: "
-                f"{len(encoder_list)} vs {len(decoder_rows)}."
-            )
-
-        device = next(self.score_model.parameters()).device
-        decoder_batch, decoder_mask = self._pad_decoder_rows(decoder_rows, device)
-        encoder_batch, encoder_mask = self._pad_encoder_outputs(encoder_list, device)
-        with torch.inference_mode():
-            outputs = self.score_model(
-                encoder_outputs=BaseModelOutput(last_hidden_state=encoder_batch),
-                attention_mask=encoder_mask,
-                decoder_input_ids=decoder_batch,
-                decoder_attention_mask=decoder_mask,
-            )
-            margins = outputs.logits.squeeze(-1).to(hidden.device, dtype=hidden.dtype)
-            _debug(f"margins={margins.float().tolist()}")
-        for row_index, last_index in enumerate(last_indices):
-            hidden[last_index, 0] = margins[row_index]
-        return hidden
-
-
-__all__ = ["T5Gemma2VllmScoreClassification"]

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/modeling_score.py

@@ -1,189 +0,0 @@
-from __future__ import annotations
-
-from typing import Optional, Tuple
-
-import torch
-import torch.nn.functional as F
-from transformers.modeling_outputs import BaseModelOutput, SequenceClassifierOutput
-from transformers.models.t5gemma2.modeling_t5gemma2 import (
-    T5Gemma2ForConditionalGeneration,
-)
-
-from .constants import (
-    DEFAULT_ENCODER_CHUNK_SIZE,
-    NO_TOKEN_ID,
-    YES_TOKEN_ID,
-)
-
-
-class T5Gemma2ForScoreClassification(T5Gemma2ForConditionalGeneration):
-    """Expose the pretrained yes/no decision as one classification logit."""
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = 1
-        self.yes_token_id = int(getattr(config, "yes_token_id", YES_TOKEN_ID))
-        self.no_token_id = int(getattr(config, "no_token_id", NO_TOKEN_ID))
-        self.encoder_chunk_size = getattr(
-            config, "encoder_chunk_size", DEFAULT_ENCODER_CHUNK_SIZE
-        )
-        if self.encoder_chunk_size is not None:
-            self.encoder_chunk_size = int(self.encoder_chunk_size)
-        self._validate_score_config()
-
-    def _validate_score_config(self) -> None:
-        vocab_size = int(getattr(self.config, "vocab_size", 0))
-        if vocab_size <= 0:
-            raise ValueError("config.vocab_size must be a positive integer.")
-        for name, token_id in (
-            ("yes_token_id", self.yes_token_id),
-            ("no_token_id", self.no_token_id),
-        ):
-            if token_id < 0 or token_id >= vocab_size:
-                raise ValueError(
-                    f"{name}={token_id} is outside vocab_size={vocab_size}."
-                )
-        if self.encoder_chunk_size is not None and self.encoder_chunk_size <= 0:
-            raise ValueError("encoder_chunk_size must be positive or None.")
-
-        self.config.num_labels = 1
-        self.config.yes_token_id = self.yes_token_id
-        self.config.no_token_id = self.no_token_id
-        self.config.encoder_chunk_size = self.encoder_chunk_size
-
-    @staticmethod
-    def _pool_encoder_chunks(
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        chunk_size: int,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        batch_size, sequence_length, hidden_size = hidden_states.shape
-        num_chunks = (sequence_length + chunk_size - 1) // chunk_size
-        padded_length = num_chunks * chunk_size
-        pad_length = padded_length - sequence_length
-        if pad_length:
-            hidden_states = F.pad(hidden_states, (0, 0, 0, pad_length))
-            attention_mask = F.pad(attention_mask, (0, pad_length))
-
-        hidden_states = hidden_states.view(
-            batch_size, num_chunks, chunk_size, hidden_size
-        )
-        chunk_mask = attention_mask.view(batch_size, num_chunks, chunk_size)
-        expanded_mask = chunk_mask.unsqueeze(-1).to(hidden_states.dtype)
-        pooled_hidden = (hidden_states * expanded_mask).sum(dim=2)
-        pooled_hidden /= chunk_mask.sum(dim=2).clamp(min=1).unsqueeze(-1)
-        pooled_mask = (chunk_mask.sum(dim=2) > 0).to(attention_mask.dtype)
-        return pooled_hidden, pooled_mask
-
-    def _forward_with_optional_chunk_pooling(
-        self,
-        *,
-        input_ids: Optional[torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        decoder_input_ids: Optional[torch.Tensor],
-        decoder_attention_mask: Optional[torch.Tensor],
-        encoder_outputs=None,
-        **kwargs,
-    ):
-        if self.encoder_chunk_size is None or encoder_outputs is not None:
-            return super().forward(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                decoder_input_ids=decoder_input_ids,
-                decoder_attention_mask=decoder_attention_mask,
-                encoder_outputs=encoder_outputs,
-                return_dict=True,
-                **kwargs,
-            )
-        if input_ids is None:
-            raise ValueError("input_ids are required when encoder_outputs is None.")
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
-
-        raw_encoder_outputs = self.get_encoder()(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            return_dict=True,
-        )
-        pooled_hidden, pooled_mask = self._pool_encoder_chunks(
-            raw_encoder_outputs.last_hidden_state,
-            attention_mask,
-            self.encoder_chunk_size,
-        )
-        return super().forward(
-            encoder_outputs=BaseModelOutput(last_hidden_state=pooled_hidden),
-            attention_mask=pooled_mask,
-            decoder_input_ids=decoder_input_ids,
-            decoder_attention_mask=decoder_attention_mask,
-            return_dict=True,
-            **kwargs,
-        )
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        decoder_input_ids: Optional[torch.Tensor] = None,
-        decoder_attention_mask: Optional[torch.Tensor] = None,
-        encoder_outputs=None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs,
-    ) -> SequenceClassifierOutput:
-        if labels is not None:
-            raise NotImplementedError("This wrapper is inference-only.")
-        if decoder_input_ids is None:
-            raise ValueError("decoder_input_ids are required.")
-        if decoder_input_ids.shape[0] == 0:
-            raise ValueError("empty batches are not supported.")
-        if input_ids is not None and input_ids.shape[0] == 0:
-            raise ValueError("empty batches are not supported.")
-
-        outputs = self._forward_with_optional_chunk_pooling(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            decoder_input_ids=decoder_input_ids,
-            decoder_attention_mask=decoder_attention_mask,
-            encoder_outputs=encoder_outputs,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            **kwargs,
-        )
-        logits = outputs.logits
-        if decoder_attention_mask is None:
-            sequence_lengths = torch.full(
-                (logits.shape[0],),
-                logits.shape[1] - 1,
-                dtype=torch.long,
-                device=logits.device,
-            )
-        else:
-            sequence_lengths = decoder_attention_mask.sum(dim=1).to(torch.long) - 1
-            if (sequence_lengths < 0).any():
-                raise ValueError("decoder_attention_mask contains an empty sequence.")
-
-        batch_indices = torch.arange(logits.shape[0], device=logits.device)
-        last_logits = logits[batch_indices, sequence_lengths]
-        yes_no_logits = torch.stack(
-            (
-                last_logits[:, self.yes_token_id],
-                last_logits[:, self.no_token_id],
-            ),
-            dim=-1,
-        ).float()
-        if not torch.isfinite(yes_no_logits).all():
-            bad_count = (~torch.isfinite(yes_no_logits).all(dim=-1)).sum().item()
-            raise RuntimeError(f"Non-finite yes/no logits for {bad_count} input(s).")
-
-        margin = yes_no_logits[:, 0] - yes_no_logits[:, 1]
-        return SequenceClassifierOutput(
-            logits=margin[:, None],
-            hidden_states=getattr(outputs, "hidden_states", None),
-            attentions=getattr(outputs, "attentions", None),
-        )
-
-
-__all__ = ["T5Gemma2ForScoreClassification"]
-

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/processing.py

@@ -1,173 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Mapping
-import hashlib
-from typing import Any
-
-import torch
-from transformers.feature_extraction_utils import BatchFeature
-from vllm.inputs import MultiModalDataDict, mm_enc_dec_input, mm_input
-from vllm.multimodal.inputs import MultiModalFieldConfig, PlaceholderRange
-from vllm.multimodal.parse import ModalityDataItems, MultiModalDataItems
-from vllm.multimodal.processing import (
-    BaseDummyInputsBuilder,
-    BaseProcessingInfo,
-    EncDecMultiModalProcessor,
-    ProcessorInputs,
-    TimingContext,
-)
-
-from .constants import TEXT_MODALITY
-
-
-class TextTokenItems(ModalityDataItems[list[str], str]):
-    def __init__(self, data: list[str]) -> None:
-        super().__init__(data, TEXT_MODALITY)
-
-    def get_count(self) -> int:
-        return len(self.data)
-
-    def get(self, index: int) -> str:
-        return self.data[index]
-
-    def get_processor_data(self) -> Mapping[str, object]:
-        return {}
-
-    def get_passthrough_data(self) -> Mapping[str, object]:
-        return {}
-
-
-class TextEncoderProcessingInfo(BaseProcessingInfo):
-    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
-        return {TEXT_MODALITY: 1}
-
-    def get_mm_max_tokens_per_item(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-    ) -> Mapping[str, int] | None:
-        return {TEXT_MODALITY: seq_len}
-
-    def parse_mm_data(
-        self,
-        mm_data: MultiModalDataDict,
-        *,
-        validate: bool = True,
-    ) -> MultiModalDataItems:
-        text_data = mm_data.get(TEXT_MODALITY)
-        if text_data is None:
-            items = TextTokenItems([])
-        elif isinstance(text_data, str):
-            items = TextTokenItems([text_data])
-        elif isinstance(text_data, list):
-            items = TextTokenItems([str(item) for item in text_data])
-        else:
-            items = TextTokenItems([str(text_data)])
-        if validate:
-            self.validate_num_items(TEXT_MODALITY, items.get_count())
-        return MultiModalDataItems({TEXT_MODALITY: items})
-
-
-class TextEncoderDummyInputsBuilder(
-    BaseDummyInputsBuilder[TextEncoderProcessingInfo]
-):
-    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
-        return "<Document>: dummy"
-
-    def get_dummy_mm_data(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, Any],
-    ) -> MultiModalDataDict:
-        count = mm_counts.get(TEXT_MODALITY, 0)
-        return {TEXT_MODALITY: ["<Document>: dummy"] * count}
-
-
-class TextEncoderProcessor(EncDecMultiModalProcessor[TextEncoderProcessingInfo]):
-    skip_decoder_start_token = True
-
-    def create_encoder_prompt(
-        self,
-        prompt: str | list[int],
-        mm_items: MultiModalDataItems,
-    ) -> str | list[int]:
-        return prompt
-
-    def _get_mm_fields_config(
-        self,
-        hf_inputs: BatchFeature,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> Mapping[str, MultiModalFieldConfig]:
-        return {
-            "encoder_input_ids": MultiModalFieldConfig.batched(
-                TEXT_MODALITY,
-                keep_on_cpu=True,
-            )
-        }
-
-    def _get_prompt_updates(self, *args, **kwargs):
-        return []
-
-    def apply(
-        self,
-        inputs: ProcessorInputs,
-        timing_ctx: TimingContext,
-    ):
-        tokenizer = self.info.get_tokenizer()
-        if isinstance(inputs.prompt, str):
-            encoder_ids = tokenizer.encode(inputs.prompt, add_special_tokens=False)
-            encoder_prompt_text = inputs.prompt
-        else:
-            encoder_ids = list(inputs.prompt)
-            encoder_prompt_text = None
-        if not encoder_ids:
-            raise ValueError("The text encoder prompt cannot be empty.")
-
-        tensor = torch.tensor([encoder_ids], dtype=torch.long)
-        mm_kwargs = self._build_mm_kwargs(tensor)
-        text_items = inputs.mm_data_items.get(TEXT_MODALITY)
-        text_for_hash = (
-            text_items.get(0)
-            if text_items is not None and text_items.get_count() > 0
-            else repr(encoder_ids)
-        )
-        digest = hashlib.sha256(str(text_for_hash).encode("utf-8")).hexdigest()
-        mm_hashes = {TEXT_MODALITY: [f"{TEXT_MODALITY}:{digest}"]}
-        mm_placeholders = {
-            TEXT_MODALITY: [PlaceholderRange(offset=0, length=len(encoder_ids))]
-        }
-        encoder_inputs = mm_input(
-            prompt_token_ids=encoder_ids,
-            prompt=encoder_prompt_text,
-            mm_kwargs=mm_kwargs,
-            mm_hashes=mm_hashes,
-            mm_placeholders=mm_placeholders,
-        )
-        return mm_enc_dec_input(
-            encoder_inputs,
-            decoder_prompt_token_ids=encoder_ids,
-            decoder_prompt=encoder_prompt_text,
-        )
-
-    def _build_mm_kwargs(self, encoder_input_ids: torch.Tensor):
-        return self._kwargs_from_batch_feature(
-            BatchFeature({"encoder_input_ids": encoder_input_ids})
-        )
-
-    def _kwargs_from_batch_feature(self, batch: BatchFeature):
-        from vllm.multimodal.inputs import MultiModalKwargsItems
-
-        return MultiModalKwargsItems.from_hf_inputs(
-            batch,
-            self._get_mm_fields_config(batch, {}),
-        )
-
-
-__all__ = [
-    "TextEncoderDummyInputsBuilder",
-    "TextEncoderProcessingInfo",
-    "TextEncoderProcessor",
-    "TextTokenItems",
-]
-

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/register.py

@@ -1,15 +0,0 @@
-from __future__ import annotations
-
-from vllm import ModelRegistry
-
-from .constants import ARCHITECTURE
-
-
-def register() -> None:
-    ModelRegistry.register_model(
-        ARCHITECTURE,
-        "kalm_t5gemma2_vllm_plugin.modeling:T5Gemma2VllmScoreClassification",
-    )
-
-
-__all__ = ["register"]

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/reranker.py

@@ -1,357 +0,0 @@
-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",
-]

vllm_support/build/lib/kalm_t5gemma2_vllm_plugin/server.py

@@ -1,215 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import threading
-import time
-from pathlib import Path
-from typing import Any, Optional
-
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, Field
-
-from .constants import (
-    DEFAULT_DOCUMENT_MAX_LENGTH,
-    DEFAULT_ENCODER_CHUNK_SIZE,
-    DEFAULT_MAX_MODEL_LEN,
-    DEFAULT_QUERY_MAX_LENGTH,
-    MODEL_ID,
-    SUPPORTED_ENCODER_CHUNK_SIZES,
-    parse_encoder_chunk_size,
-)
-from .reranker import KaLMVLLMReranker
-
-
-class RerankRequest(BaseModel):
-    query: str
-    documents: list[str] = Field(min_length=1)
-    instruction: Optional[str] = None
-    top_k: Optional[int] = None
-    return_margin: bool = False
-
-
-class ScorePair(BaseModel):
-    query: str
-    document: str
-    id: Optional[Any] = None
-
-
-class ScoreRequest(BaseModel):
-    pairs: list[ScorePair] = Field(min_length=1)
-    instruction: Optional[str] = None
-    return_margin: bool = False
-
-
-class OnlineRerankerService:
-    def __init__(self, args: argparse.Namespace) -> None:
-        self.model = str(args.model)
-        self.query_max_length = int(args.query_max_length)
-        self.document_max_length = int(args.document_max_length)
-        self.encoder_chunk_size = parse_encoder_chunk_size(args.encoder_chunk_size)
-        self.max_model_len = int(args.max_model_len)
-        self.batch_size = int(args.batch_size)
-        self.dtype = str(args.dtype)
-        self.gpu_memory_utilization = float(args.gpu_memory_utilization)
-        self.started_at = time.time()
-        self.lock = threading.Lock()
-        self.reranker = KaLMVLLMReranker(
-            self.model,
-            query_max_length=self.query_max_length,
-            document_max_length=self.document_max_length,
-            encoder_chunk_size=self.encoder_chunk_size,
-            max_model_len=self.max_model_len,
-            batch_size=self.batch_size,
-            dtype=self.dtype,
-            gpu_memory_utilization=self.gpu_memory_utilization,
-            tensor_parallel_size=args.tensor_parallel_size,
-        )
-
-    def config(self) -> dict[str, Any]:
-        return {
-            "model": self.model,
-            "query_max_length": self.query_max_length,
-            "document_max_length": self.document_max_length,
-            "encoder_chunk_size": self.encoder_chunk_size,
-            "max_model_len": self.max_model_len,
-            "batch_size": self.batch_size,
-            "dtype": self.dtype,
-            "gpu_memory_utilization": self.gpu_memory_utilization,
-            "tensor_parallel_size": 1,
-            "supported_encoder_chunk_sizes": sorted(
-                SUPPORTED_ENCODER_CHUNK_SIZES
-            ),
-        }
-
-    def health(self) -> dict[str, Any]:
-        return {
-            "status": "ok",
-            "uptime_seconds": round(time.time() - self.started_at, 3),
-            **self.config(),
-        }
-
-    def close(self) -> None:
-        self.reranker.close()
-
-    def rerank(self, request: RerankRequest) -> dict[str, Any]:
-        if request.top_k is not None and request.top_k < 0:
-            raise ValueError("top_k must be non-negative or null.")
-        with self.lock:
-            rankings = self.reranker.rank(
-                request.query,
-                request.documents,
-                instruction=request.instruction,
-                top_k=request.top_k,
-                return_margin=request.return_margin,
-            )
-        results: list[dict[str, Any]] = []
-        for item in rankings:
-            result = {
-                "index": int(item["corpus_id"]),
-                "score": float(item["score"]),
-            }
-            if request.return_margin:
-                result["margin"] = float(item["margin"])
-            results.append(result)
-        return {"object": "rerank", "results": results}
-
-    def score(self, request: ScoreRequest) -> dict[str, Any]:
-        pairs = [(item.query, item.document) for item in request.pairs]
-        with self.lock:
-            predictions = self.reranker.predict(
-                pairs,
-                instruction=request.instruction,
-                return_margin=True,
-            )
-        results: list[dict[str, Any]] = []
-        for index, (pair, prediction) in enumerate(zip(request.pairs, predictions)):
-            assert isinstance(prediction, dict)
-            result = {
-                "index": index,
-                "score": float(prediction["score"]),
-            }
-            if pair.id is not None:
-                result["id"] = pair.id
-            if request.return_margin:
-                result["margin"] = float(prediction["margin"])
-            results.append(result)
-        return {"object": "score", "results": results}
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        description="FastAPI server for the KaLM vLLM adapter.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    parser.add_argument("--host", default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--model", default=MODEL_ID)
-    parser.add_argument("--query-max-length", type=int, default=DEFAULT_QUERY_MAX_LENGTH)
-    parser.add_argument(
-        "--document-max-length", type=int, default=DEFAULT_DOCUMENT_MAX_LENGTH
-    )
-    parser.add_argument(
-        "--encoder-chunk-size", default=str(DEFAULT_ENCODER_CHUNK_SIZE)
-    )
-    parser.add_argument("--max-model-len", type=int, default=DEFAULT_MAX_MODEL_LEN)
-    parser.add_argument("--batch-size", type=int, default=32)
-    parser.add_argument("--dtype", default="bfloat16")
-    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85)
-    parser.add_argument("--tensor-parallel-size", type=int, default=1)
-    parser.add_argument("--log-level", default="info")
-    return parser
-
-
-def create_app(service: OnlineRerankerService) -> FastAPI:
-    app = FastAPI(title="KaLM vLLM Online Reranker", version="0.1.0")
-    app.state.service = service
-    app.add_event_handler("shutdown", service.close)
-
-    @app.get("/health")
-    def health():
-        return app.state.service.health()
-
-    @app.post("/rerank")
-    def rerank(request: RerankRequest):
-        try:
-            return app.state.service.rerank(request)
-        except (TypeError, ValueError) as error:
-            raise HTTPException(status_code=400, detail=str(error)) from error
-
-    @app.post("/score")
-    def score(request: ScoreRequest):
-        try:
-            return app.state.service.score(request)
-        except (TypeError, ValueError) as error:
-            raise HTTPException(status_code=400, detail=str(error)) from error
-
-    return app
-
-
-def main() -> int:
-    args = build_parser().parse_args()
-    args.encoder_chunk_size = parse_encoder_chunk_size(args.encoder_chunk_size)
-    print("=== KaLM vLLM Online Reranker ===", flush=True)
-    print(f"model: {args.model}", flush=True)
-    print(f"query_max_length: {args.query_max_length}", flush=True)
-    print(f"document_max_length: {args.document_max_length}", flush=True)
-    print(f"encoder_chunk_size: {args.encoder_chunk_size}", flush=True)
-    print(f"max_model_len: {args.max_model_len}", flush=True)
-    print(f"listen: http://{args.host}:{args.port}", flush=True)
-
-    service = OnlineRerankerService(args)
-    app = create_app(service)
-    import uvicorn
-
-    uvicorn.run(
-        app,
-        host=args.host,
-        port=args.port,
-        log_level=args.log_level,
-        workers=1,
-    )
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

vllm_support/src/kalm_t5gemma2_vllm_plugin.egg-info/PKG-INFO

@@ -1,292 +0,0 @@
-Metadata-Version: 2.4
-Name: kalm-t5gemma2-vllm-plugin
-Version: 0.1.0
-Summary: vLLM 0.19.1 adapter for KaLM-Reranker-V1-Nano
-License: Apache-2.0
-Requires-Python: <3.13,>=3.12
-Description-Content-Type: text/markdown
-Requires-Dist: transformers==5.6.2
-Requires-Dist: vllm==0.19.1
-Provides-Extra: online
-Requires-Dist: fastapi<0.137,>=0.136; extra == "online"
-Requires-Dist: uvicorn<0.47,>=0.46; extra == "online"
-
-# KaLM-Reranker-V1-Nano vLLM Support
-
-This directory contains the experimental vLLM 0.19.1 adapter for
-`KaLM-Embedding/KaLM-Reranker-V1-Nano`. It supports offline Python and CLI
-reranking plus an optional FastAPI service.
-
-The adapter does not modify or retrain the checkpoint. It reads the original
-decoder logits for the single-token answers `yes` and `no` and returns:
-
-```text
-margin = yes_logit - no_logit
-score = sigmoid(margin) = P(yes)
-```
-
-## Tested environment
-
-- Linux and NVIDIA CUDA
-- Python 3.12
-- vLLM 0.19.1
-- Transformers 5.6.2
-- PyTorch 2.10.0
-- BF16, one GPU
-
-The package intentionally rejects other vLLM versions and
-`tensor_parallel_size != 1`. These combinations have not been validated.
-
-## Installation
-
-Create an environment and download the model repository:
-
-```bash
-conda create -n kalm-vllm python=3.12 -y
-conda activate kalm-vllm
-
-pip install "vllm==0.19.1" "transformers==5.6.2"
-pip install "fastapi>=0.136,<0.137" "uvicorn>=0.46,<0.47"
-
-hf download KaLM-Embedding/KaLM-Reranker-V1-Nano \
-  --local-dir ./KaLM-Reranker-V1-Nano
-pip install ./KaLM-Reranker-V1-Nano/vllm_support --no-deps
-export VLLM_PLUGINS=kalm_t5gemma2
-```
-
-The model can also be loaded directly by its Hugging Face ID. In that case,
-only download the `vllm_support` directory before installing the plugin:
-
-```bash
-hf download KaLM-Embedding/KaLM-Reranker-V1-Nano \
-  --include "vllm_support/**" \
-  --local-dir ./KaLM-Reranker-V1-Nano
-pip install ./KaLM-Reranker-V1-Nano/vllm_support --no-deps
-```
-
-## Offline Python API
-
-```python
-from kalm_t5gemma2_vllm_plugin import KaLMVLLMReranker
-
-query = "What is the capital of China?"
-documents = [
-    "The capital of China is Beijing.",
-    "Gravity attracts bodies toward one another.",
-]
-
-pairs = [(query, document) for document in documents]
-with KaLMVLLMReranker(
-    "KaLM-Embedding/KaLM-Reranker-V1-Nano",
-    query_max_length=512,
-    document_max_length=1024,
-    encoder_chunk_size=4,
-    max_model_len=2048,
-    batch_size=32,
-) as reranker:
-    print(reranker.predict(pairs))
-    print(reranker.predict(pairs, return_margin=True))
-    print(reranker.rank(query, documents))
-```
-
-Expected BF16 scores are approximately:
-
-```text
-[0.99980897, 0.00000493699]
-```
-
-`predict()` preserves input order. `rank()` returns score-descending results
-with the original document index in `corpus_id`.
-
-## Offline CLI
-
-Run the built-in example:
-
-```bash
-kalm-vllm-rerank --return-margin
-```
-
-Score JSONL input:
-
-```bash
-kalm-vllm-rerank \
-  --input-jsonl ./KaLM-Reranker-V1-Nano/vllm_support/examples/sample_pairs.jsonl \
-  --output-jsonl ./scores.jsonl \
-  --return-margin
-```
-
-Each input line must contain `query` and `document`. Optional fields are `id`
-and `instruction`. `--top-k N` groups rows by exact query text, sorts each
-group by score, and keeps its first `N` documents.
-
-## Online service
-
-Start one model instance:
-
-```bash
-kalm-vllm-serve \
-  --host 0.0.0.0 \
-  --port 8000 \
-  --model KaLM-Embedding/KaLM-Reranker-V1-Nano \
-  --encoder-chunk-size 4
-```
-
-The portable startup script exposes the same settings through environment
-variables:
-
-```bash
-CUDA_VISIBLE_DEVICES=0 PORT=8000 \
-  ./KaLM-Reranker-V1-Nano/vllm_support/examples/start_online_server.sh
-```
-
-In a second terminal, check health and send built-in demo requests:
-
-```bash
-kalm-vllm-client --health
-kalm-vllm-client --endpoint rerank --return-margin
-kalm-vllm-client --endpoint score --return-margin
-```
-
-For custom input, pass one JSON object with `--json-file`. Use `/rerank` for
-one query against multiple documents:
-
-```bash
-kalm-vllm-client \
-  --endpoint rerank \
-  --json-file ./KaLM-Reranker-V1-Nano/vllm_support/examples/rerank_request.json \
-  --return-margin \
-  --top-k 10
-```
-
-Use `/score` for a batch of independent query-document pairs:
-
-```bash
-kalm-vllm-client \
-  --endpoint score \
-  --json-file ./KaLM-Reranker-V1-Nano/vllm_support/examples/score_request.json \
-  --return-margin
-```
-
-When `--json-file` is used, `--return-margin` sets
-`"return_margin": true` in the outgoing request, and `--top-k` overrides the
-JSON value for `/rerank`.
-
-### `POST /rerank`
-
-```json
-{
-  "query": "What is the capital of China?",
-  "documents": [
-    "The capital of China is Beijing.",
-    "Gravity attracts bodies toward one another."
-  ],
-  "instruction": "Given a query, retrieve documents that answer the query.",
-  "top_k": null,
-  "return_margin": true
-}
-```
-
-Results are returned in descending score order:
-
-```json
-{
-  "object": "rerank",
-  "results": [
-    {"index": 0, "score": 0.9998089, "margin": 8.5625},
-    {"index": 1, "score": 0.00000493699, "margin": -12.21875}
-  ]
-}
-```
-
-### `POST /score`
-
-```json
-{
-  "pairs": [
-    {
-      "id": "doc-1",
-      "query": "What is the capital of China?",
-      "document": "The capital of China is Beijing."
-    }
-  ],
-  "instruction": null,
-  "return_margin": false
-}
-```
-
-`/score` accepts multiple entries in `pairs`, preserves their input order and
-includes an input `id` when provided.
-
-### `GET /health`
-
-Returns service status and the effective model, length, chunking, dtype and
-memory settings.
-
-## Configuration
-
-| Setting | Default | Meaning |
-| --- | ---: | --- |
-| `query_max_length` | `512` | Maximum raw query tokens before prompt insertion |
-| `document_max_length` | `1024` | Maximum encoder tokens for `<Document>: ...` |
-| `encoder_chunk_size` | `4` | Mean-pooling factor; one of `1,2,4,8,16,32` |
-| `max_model_len` | `2048` | vLLM engine context budget |
-| `batch_size` | `32` | Pairs passed to each `LLM.classify()` call |
-| `dtype` | `bfloat16` | Model compute dtype |
-| `gpu_memory_utilization` | `0.85` | vLLM GPU memory fraction |
-| `tensor_parallel_size` | `1` | Only supported value in this release |
-
-The query and document limits belong to separate decoder and encoder streams;
-they are not a combined cross-encoder token limit. Larger values are
-configurable but have not been validated up to the model card's full 128K
-limit.
-
-## Limitations
-
-- This is a custom `LLM.classify()` plugin, not vLLM's native HTTP `/score`
-  implementation.
-- The shim uses vLLM scheduling and pooling interfaces but executes the
-  T5Gemma2 semantic forward through Transformers. It is not a complete
-  vLLM-native kernel implementation and should not be used to claim native
-  vLLM throughput.
-- Online serving is a single-process FastAPI wrapper around one model instance.
-- `encoder_chunk_size=None`, `null`, or an empty string falls back to `4`; it
-  does not disable pooling in this release.
-
-## Troubleshooting
-
-**The plugin is not discovered**
-
-Reinstall the package and ensure the environment variable includes its entry
-point name:
-
-```bash
-pip install ./KaLM-Reranker-V1-Nano/vllm_support --no-deps --force-reinstall
-export VLLM_PLUGINS=kalm_t5gemma2
-```
-
-**The adapter reports an unsupported vLLM version**
-
-Install exactly `vllm==0.19.1`. Internal model and processor APIs are version
-sensitive.
-
-**The tokenizer check fails**
-
-Confirm that the tokenizer belongs to this Nano checkpoint. The adapter
-requires `yes -> 4443` and `no -> 1904`.
-
-**CUDA runs out of memory**
-
-Reduce `batch_size`, `document_max_length`, `query_max_length`,
-`max_model_len`, or `gpu_memory_utilization`.
-
-**CUDA initialization fails with error 803**
-
-The process may be resolving a CUDA compatibility library before the host
-driver library. On common Debian/Ubuntu layouts, retry with:
-
-```bash
-export LD_LIBRARY_PATH="/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
-```
-
-The provided `start_online_server.sh` applies this adjustment automatically
-when both directories exist.

vllm_support/src/kalm_t5gemma2_vllm_plugin.egg-info/SOURCES.txt

@@ -1,18 +0,0 @@
-README.md
-pyproject.toml
-src/kalm_t5gemma2_vllm_plugin/__init__.py
-src/kalm_t5gemma2_vllm_plugin/cli.py
-src/kalm_t5gemma2_vllm_plugin/client.py
-src/kalm_t5gemma2_vllm_plugin/constants.py
-src/kalm_t5gemma2_vllm_plugin/modeling.py
-src/kalm_t5gemma2_vllm_plugin/modeling_score.py
-src/kalm_t5gemma2_vllm_plugin/processing.py
-src/kalm_t5gemma2_vllm_plugin/register.py
-src/kalm_t5gemma2_vllm_plugin/reranker.py
-src/kalm_t5gemma2_vllm_plugin/server.py
-src/kalm_t5gemma2_vllm_plugin.egg-info/PKG-INFO
-src/kalm_t5gemma2_vllm_plugin.egg-info/SOURCES.txt
-src/kalm_t5gemma2_vllm_plugin.egg-info/dependency_links.txt
-src/kalm_t5gemma2_vllm_plugin.egg-info/entry_points.txt
-src/kalm_t5gemma2_vllm_plugin.egg-info/requires.txt
-src/kalm_t5gemma2_vllm_plugin.egg-info/top_level.txt

vllm_support/src/kalm_t5gemma2_vllm_plugin.egg-info/dependency_links.txt

@@ -1 +0,0 @@
-

vllm_support/src/kalm_t5gemma2_vllm_plugin.egg-info/entry_points.txt

@@ -1,7 +0,0 @@
-[console_scripts]
-kalm-vllm-client = kalm_t5gemma2_vllm_plugin.client:main
-kalm-vllm-rerank = kalm_t5gemma2_vllm_plugin.cli:main
-kalm-vllm-serve = kalm_t5gemma2_vllm_plugin.server:main
-
-[vllm.general_plugins]
-kalm_t5gemma2 = kalm_t5gemma2_vllm_plugin.register:register

vllm_support/src/kalm_t5gemma2_vllm_plugin.egg-info/requires.txt

@@ -1,6 +0,0 @@
-transformers==5.6.2
-vllm==0.19.1
-
-[online]
-fastapi<0.137,>=0.136
-uvicorn<0.47,>=0.46

vllm_support/src/kalm_t5gemma2_vllm_plugin.egg-info/top_level.txt

@@ -1 +0,0 @@
-kalm_t5gemma2_vllm_plugin

vllm_support/src/kalm_t5gemma2_vllm_plugin/client.py

@@ -102,7 +102,7 @@ def _request_payload(args: argparse.Namespace) -> dict[str, Any]:
         payload["top_k"] = args.top_k
     return payload
 
-
+    
 def main() -> int:
     args = build_parser().parse_args()
     base_url = args.base_url.rstrip("/")

vllm_support/src/kalm_t5gemma2_vllm_plugin/server.py

@@ -163,7 +163,7 @@ def build_parser() -> argparse.ArgumentParser:
 def create_app(service: OnlineRerankerService) -> FastAPI:
     app = FastAPI(title="KaLM vLLM Online Reranker", version="0.1.0")
     app.state.service = service
-    app.add_event_handler("shutdown", service.close)
+    app.router.add_event_handler("shutdown", service.close)
 
     @app.get("/health")
     def health():