Is this page helpful?

API Usage Examples (gRPC) for NVIDIA NeMo Retriever Reranking NIM#

Use the examples in this documentation to help you get started using the API for NVIDIA NeMo Retriever Reranking NIM.

For the full API reference, refer to API Reference (gRPC) for NVIDIA NeMo Retriever Reranking NIM.

gRPC Support#

The NeMo Retriever Reranking NIM supports the Open Inference Protocol (KServe V2). You can make gRPC inference requests by using the KServe V2 protocol buffers with a gRPC client.

Launch the NeMo Retriever Reranking NIM#

Launch NeMo Retriever Reranking NIM by following the Get Started guide. In the code to launch the NIM, include -e NIM_SERVER_GRPC_BIND_ADDR=0.0.0.0:8001 to enable the gRPC listener and -p 8001:8001 to publish its port, as shown in the following example.

# Start the NIM with the gRPC listener enabled and published.
# The NIM_ENGINE_* lines set the model, download provider, model path, and precision explicitly.
# Keep them to make the settings explicit, or delete them to use the defaults.
docker run -it --rm --name=$CONTAINER_NAME \
  --runtime=nvidia \
  --gpus all \
  --shm-size=16GB \
  -e NIM_ENGINE_MODEL_DOWNLOAD_PROVIDER=hf \
  -e HF_TOKEN \
  -e NIM_ENGINE_MODEL_NAME=nvidia/Nemotron-3.5-Rerank-8B-BF16 \
  -e NIM_ENGINE_MODEL_PATH=/model/rerank \
  -e NIM_ENGINE_PRECISION=fp16 \
  -e NIM_SERVER_GRPC_BIND_ADDR=0.0.0.0:8001 \
  -v "$LOCAL_NIM_CACHE/cache:/opt/cache" \
  -v "$LOCAL_NIM_CACHE/weights:/model" \
  -u $(id -u) \
  -p 8000:8000 \
  -p 8001:8001 \
  $IMG_NAME

HTTP remains available on port 8000. Both protocols use the same model and worker pool. The gRPC listener uses plaintext transport; HTTP TLS settings don’t configure gRPC TLS.

Make Inference Calls#

After you launch the NeMo Retriever Reranking NIM, you can make inference calls by using the following code.

First, install Python dependencies.

python3 -m pip install grpcio grpcio-tools numpy
curl -LO https://raw.githubusercontent.com/kserve/open-inference-protocol/main/specification/protocol/open_inference_grpc.proto
python3 -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. open_inference_grpc.proto

Next, make inference calls.

import asyncio
from dataclasses import dataclass
from enum import StrEnum, auto

import grpc
import numpy as np
from numpy.typing import NDArray

import open_inference_grpc_pb2 as oip
import open_inference_grpc_pb2_grpc as oip_grpc


class InputName(StrEnum):
    QUERY = auto()
    PASSAGE = auto()


class TruncateValue(StrEnum):
    END = "END"
    NONE = "NONE"


class OutputName(StrEnum):
    INDEX = auto()
    LOGIT = auto()
    TOKEN_COUNT = auto()


@dataclass
class RankingResponse:
    index: NDArray[np.int32]
    logit: NDArray[np.float32]
    token_count: NDArray[np.int32]


def _string_param(value: StrEnum | str) -> oip.InferParameter:
    return oip.InferParameter(string_param=value.value if isinstance(value, StrEnum) else value)


def _bytes_tensor(name: str, values: list[str]) -> oip.ModelInferRequest.InferInputTensor:
    encoded_values = [value.encode("utf-8") for value in values]
    return oip.ModelInferRequest.InferInputTensor(
        name=name,
        datatype="BYTES",
        shape=[len(encoded_values), 1],
        contents=oip.InferTensorContents(bytes_contents=encoded_values),
    )


def _output_by_name(
    response: oip.ModelInferResponse, name: OutputName
) -> tuple[int, oip.ModelInferResponse.InferOutputTensor]:
    for index, output in enumerate(response.outputs):
        if output.name == name.value:
            return index, output
    raise ValueError(f"'{name.value}' not found in the response")


def _int32_output(response: oip.ModelInferResponse, name: OutputName) -> NDArray[np.int32]:
    index, output = _output_by_name(response, name)
    if index < len(response.raw_output_contents):
        return np.frombuffer(response.raw_output_contents[index], dtype="<i4").reshape(
            output.shape
        )
    return np.array(output.contents.int_contents, dtype=np.int32).reshape(output.shape)


def _fp32_output(response: oip.ModelInferResponse, name: OutputName) -> NDArray[np.float32]:
    index, output = _output_by_name(response, name)
    if index < len(response.raw_output_contents):
        return np.frombuffer(response.raw_output_contents[index], dtype="<f4").reshape(
            output.shape
        )
    return np.array(output.contents.fp32_contents, dtype=np.float32).reshape(output.shape)


class RankingClient:
    def __init__(self, endpoint: str = "localhost:8001") -> None:
        self._endpoint = endpoint

    async def compute(
        self,
        model_name: str,
        query: str,
        passages: list[str],
        truncate: TruncateValue = TruncateValue.NONE,
    ) -> RankingResponse:
        request = oip.ModelInferRequest(
            model_name=model_name,
            inputs=[
                _bytes_tensor(InputName.QUERY.value, [query]),
                _bytes_tensor(InputName.PASSAGE.value, passages),
            ],
            parameters={"truncate": _string_param(truncate)},
        )

        async with grpc.aio.insecure_channel(self._endpoint) as channel:
            client = oip_grpc.GRPCInferenceServiceStub(channel)
            response = await client.ModelInfer(request)

        return RankingResponse(
            index=_int32_output(response, OutputName.INDEX),
            logit=_fp32_output(response, OutputName.LOGIT),
            token_count=_int32_output(response, OutputName.TOKEN_COUNT),
        )


async def main() -> None:
    client = RankingClient()
    response = await client.compute(
        model_name="nvidia_nemotron_3_5_rerank_8b_bf16",
        query="Which animal is a mammal?",
        passages=["A dolphin is a mammal.", "A trout is a fish."],
        truncate=TruncateValue.NONE,
    )
    print(response.index.shape)
    print(response.logit.shape)
    print(response.token_count.shape)


asyncio.run(main())

The example returns one ranking, score, and token count per passage. It prints the three output shapes, as shown in the following output.

(2,)
(2,)
(2,)

All three outputs use descending score order. Each index identifies a passage in the original request. The corresponding token_count includes the effective prompt and image tokens for that query and passage pair. For more information about image passages, raw tensor encoding, output subsets, and served model aliases, refer to API Reference (gRPC) for NVIDIA NeMo Retriever Reranking NIM.