Upload 34 files

vllm_support/README.md

@@ -0,0 +1,279 @@
+# KaLM-Reranker-V1-Small 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
+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-Small \
+  --local-dir ./KaLM-Reranker-V1-Small
+pip install ./KaLM-Reranker-V1-Small/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-Small \
+  --include "vllm_support/**" \
+  --local-dir ./KaLM-Reranker-V1-Small
+pip install ./KaLM-Reranker-V1-Small/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-Small",
+    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-Small/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-Small \
+  --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-Small/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-Small/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-Small/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-Small/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 Small 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/build/lib/kalm_t5gemma2_vllm_plugin/__init__.py

@@ -0,0 +1,23 @@
+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

@@ -0,0 +1,172 @@
+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

@@ -0,0 +1,125 @@
+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

@@ -0,0 +1,105 @@
+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

@@ -0,0 +1,265 @@
+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

@@ -0,0 +1,189 @@
+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

@@ -0,0 +1,173 @@
+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

@@ -0,0 +1,15 @@
+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

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

@@ -0,0 +1,215 @@
+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/examples/call_online_api.py

@@ -0,0 +1,6 @@
+from kalm_t5gemma2_vllm_plugin.client import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

vllm_support/examples/offline_usage.py

@@ -0,0 +1,17 @@
+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.",
+]
+
+with KaLMVLLMReranker(
+    "KaLM-Embedding/KaLM-Reranker-V1-Nano",
+    encoder_chunk_size=4,
+    query_max_length=512,
+    document_max_length=1024,
+) as reranker:
+    print(reranker.predict([(query, document) for document in documents]))
+    print(reranker.rank(query, documents))

vllm_support/examples/rerank_request.json

@@ -0,0 +1,10 @@
+{
+  "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": false
+}

vllm_support/examples/sample_pairs.jsonl

@@ -0,0 +1,2 @@
+{"id":"positive","query":"What is the capital of China?","document":"The capital of China is Beijing."}
+{"id":"negative","query":"What is the capital of China?","document":"Gravity attracts bodies toward one another."}

vllm_support/examples/score_request.json

@@ -0,0 +1,16 @@
+{
+  "pairs": [
+    {
+      "id": "positive",
+      "query": "What is the capital of China?",
+      "document": "The capital of China is Beijing."
+    },
+    {
+      "id": "negative",
+      "query": "What is the capital of China?",
+      "document": "Gravity attracts bodies toward one another."
+    }
+  ],
+  "instruction": null,
+  "return_margin": false
+}

vllm_support/examples/start_online_server.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
+export VLLM_PLUGINS="${VLLM_PLUGINS:-kalm_t5gemma2}"
+if [[ -d /lib/x86_64-linux-gnu && -d /usr/lib/x86_64-linux-gnu ]]; then
+  export LD_LIBRARY_PATH="/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+fi
+
+exec kalm-vllm-serve \
+  --host "${HOST:-0.0.0.0}" \
+  --port "${PORT:-8000}" \
+  --model "${MODEL:-KaLM-Embedding/KaLM-Reranker-V1-Nano}" \
+  --query-max-length "${QUERY_MAX_LENGTH:-512}" \
+  --document-max-length "${DOCUMENT_MAX_LENGTH:-1024}" \
+  --encoder-chunk-size "${ENCODER_CHUNK_SIZE:-4}" \
+  --max-model-len "${MAX_MODEL_LEN:-2048}" \
+  --batch-size "${BATCH_SIZE:-32}" \
+  --dtype "${DTYPE:-bfloat16}" \
+  --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.85}"

vllm_support/pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kalm-t5gemma2-vllm-plugin"
+version = "0.1.0"
+description = "vLLM 0.19.1 adapter for KaLM-Reranker-V1-Nano"
+readme = "README.md"
+requires-python = ">=3.12,<3.13"
+license = {text = "Apache-2.0"}
+dependencies = [
+    "transformers==5.6.2",
+    "vllm==0.19.1",
+]
+
+[project.optional-dependencies]
+online = [
+    "fastapi>=0.136,<0.137",
+    "uvicorn>=0.46,<0.47",
+]
+
+[project.entry-points."vllm.general_plugins"]
+kalm_t5gemma2 = "kalm_t5gemma2_vllm_plugin.register:register"
+
+[project.scripts]
+kalm-vllm-rerank = "kalm_t5gemma2_vllm_plugin.cli:main"
+kalm-vllm-serve = "kalm_t5gemma2_vllm_plugin.server:main"
+kalm-vllm-client = "kalm_t5gemma2_vllm_plugin.client:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+

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

@@ -0,0 +1,292 @@
+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

@@ -0,0 +1,18 @@
+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

@@ -0,0 +1 @@
+

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

@@ -0,0 +1,7 @@
+[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

@@ -0,0 +1,6 @@
+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

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

vllm_support/src/kalm_t5gemma2_vllm_plugin/__init__.py

@@ -0,0 +1,23 @@
+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/src/kalm_t5gemma2_vllm_plugin/cli.py

@@ -0,0 +1,172 @@
+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/src/kalm_t5gemma2_vllm_plugin/client.py

@@ -0,0 +1,125 @@
+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/src/kalm_t5gemma2_vllm_plugin/constants.py

@@ -0,0 +1,105 @@
+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/src/kalm_t5gemma2_vllm_plugin/modeling.py

@@ -0,0 +1,265 @@
+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/src/kalm_t5gemma2_vllm_plugin/modeling_score.py

@@ -0,0 +1,189 @@
+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/src/kalm_t5gemma2_vllm_plugin/processing.py

@@ -0,0 +1,173 @@
+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/src/kalm_t5gemma2_vllm_plugin/register.py

@@ -0,0 +1,15 @@
+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/src/kalm_t5gemma2_vllm_plugin/reranker.py

@@ -0,0 +1,357 @@
+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/src/kalm_t5gemma2_vllm_plugin/server.py

@@ -0,0 +1,215 @@
+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())