Nvidia NIM - Rerank

Nvidia NIM - Rerank

LiteLLM에서 Nvidia NIM 리랭크(재정렬) 모델을 사용하는 방법을 알아봐요. 의미 검색과 RAG(검색 증강 생성)에 유용해요.

출처: 문서

본문

속성 내용
설명 Nvidia NIM은 의미 검색과 RAG를 위한 고성능 리랭크 모델을 제공해요
제공사 문서 Nvidia NIM Rerank API ↗
지원 엔드포인트 /rerank

개요

Nvidia NIM 리랭크 모델은 다음을 도와줘요:

  • 쿼리에 대한 관련성으로 검색 결과 재정렬
  • RAG(Retrieval-Augmented Generation) 정확도 향상
  • 대규모 문서 집합을 효율적으로 필터링·랭킹

지원 모델:

  • 플랫폼의 모든 Nvidia NIM 리랭크 모델

팁: LiteLLM이 지원하는 Nvidia NIM 리랭크 모델의 전체 목록은 Nvidia NIM에서 확인할 수 있어요.

사용법

LiteLLM Python SDK

LLaMa 1B 모델:

import litellm
import os

os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."

response = litellm.rerank(
    model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
    query="What is the GPU memory bandwidth of H100 SXM?",
    documents=[
        "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.",
        "A100 provides up to 20X higher performance over the prior generation.",
        "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
    ],
    top_n=3,
)

print(response)

Mistral 4B 모델:

import litellm
import os

os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."

response = litellm.rerank(
    model="nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3",
    query="What is the GPU memory bandwidth of H100 SXM?",
    documents=[
        "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.",
        "A100 provides up to 20X higher performance over the prior generation.",
        "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
    ],
    top_n=3,
)

print(response)

응답:

{
    "results": [
        {
            "index": 2,
            "relevance_score": 6.828125,
            "document": {
                "text": "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
            }
        },
        {
            "index": 0,
            "relevance_score": -1.564453125,
            "document": {
                "text": "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth."
            }
        }
    ]
}

LiteLLM Proxy 사용법

1. Config 설정

Nvidia NIM 리랭크 모델을 proxy 설정에 추가해요:

model_list:
  - model_name: nvidia-rerank
    litellm_params:
      model: nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2
      api_key: os.environ/NVIDIA_NIM_API_KEY

2. Proxy 시작

litellm --config /path/to/config.yaml

3. 리랭크 요청 보내기

curl -X POST http://0.0.0.0:4000/rerank \
  -H "Authorization: Bearer ***" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "nvidia-rerank",
    "query": "What is the GPU memory bandwidth of H100?",
    "documents": [
      "H100 delivers 3TB/s memory bandwidth",
      "A100 has 2TB/s memory bandwidth",
      "V100 offers 900GB/s memory bandwidth"
    ],
    "top_n": 2
  }'

/v1/ranking 모델 (llama-3.2-nv-rerankqa-1b-v2)

일부 Nvidia NIM 리랭크 모델은 기본 /v1/retrieval/{model}/reranking 엔드포인트 대신 /v1/ranking 엔드포인트를 사용해요.

ranking/ 접두사를 사용해 요청을 강제로 /v1/ranking 엔드포인트로 보낼 수 있어요:

LiteLLM Python SDK

import litellm
import os

os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."

# Use "ranking/" prefix to force /v1/ranking endpoint
response = litellm.rerank(
    model="nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2",
    query="which way did the traveler go?",
    documents=[
        "two roads diverged in a yellow wood...",
        "then took the other, as just as fair...",
        "i shall be telling this with a sigh somewhere ages and ages hence..."
    ],
    top_n=3,
    truncate="END",  # Optional: truncate long text from the end
)

print(response)

LiteLLM Proxy

config.yaml:

model_list:
  - model_name: nvidia-ranking
    litellm_params:
      model: nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2
      api_key: os.environ/NVIDIA_NIM_API_KEY

LiteLLM Proxy에 요청:

curl -X POST http://0.0.0.0:4000/rerank \
  -H "Authorization: Bearer ***" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "nvidia-ranking",
    "query": "which way did the traveler go?",
    "documents": [
      "two roads diverged in a yellow wood...",
      "then took the other, as just as fair..."
    ],
    "top_n": 2
  }'

모델 해석 이해하기

/v1/ranking 엔드포인트:

model: nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2
       └────┬────┘ └──┬──┘ └─────────────┬──────────────────┘
            │        │                   │
            │        │                   └────▶ Model name sent to provider
            │        │
            │        └────────────────────────▶ Tells LiteLLM the request/response and url should be sent to Nvidia NIM /v1/ranking endpoint
            │
            └─────────────────────────────────▶ Provider prefix

API URL: https://ai.api.nvidia.com/v1/ranking

시각적 흐름:

Client Request                LiteLLM                              Provider API
──────────────              ────────────                         ─────────────

# Default reranking endpoint
model: "nvidia_nim/nvidia/model-name"
                            1. Extracts model: nvidia/model-name
                            2. Routes to default endpoint ──────▶ POST /v1/retrieval/nvidia/model-name/reranking

# Forced ranking endpoint
model: "nvidia_nim/ranking/nvidia/model-name"
                            1. Detects "ranking/" prefix
                            2. Extracts model: nvidia/model-name
                            3. Routes to ranking endpoint ──────▶ POST /v1/ranking
                                                                  Body: {"model": "nvidia/model-name", ...}

각 엔드포인트 사용 시점:

엔드포인트 모델 접두사 사용 사례
/v1/retrieval/{model}/reranking nvidia_nim/ 대부분의 리랭크 모델의 기본값
/v1/ranking nvidia_nim/ranking/ nvidia/llama-3.2-nv-rerankqa-1b-v2처럼 이 엔드포인트를 요구하는 모델용

팁: Nvidia NIM 모델 배포 페이지에서 모델이 요구하는 엔드포인트를 확인할 수 있어요.

API 파라미터

필수 파라미터

파라미터 타입 설명
model string nvidia_nim/ 접두사가 붙은 Nvidia NIM 리랭크 모델명
query string 문서를 랭킹할 검색 쿼리
documents array 랭킹할 문서 목록 (1-1000개)

선택 파라미터

파라미터 타입 기본값 설명
top_n integer 모든 문서 반환할 상위 랭킹 문서 수

Nvidia 전용 파라미터

truncate: 텍스트가 모델의 컨텍스트 윈도우를 초과할 때 어떻게 잘라낼지 제어해요

  • "NONE": 잘라내지 않음 (너무 길면 요청이 실패할 수 있음)
  • "END": 텍스트 끝에서 잘라냄
response = litellm.rerank(
    model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
    query="GPU performance",
    documents=["High performance computing", "Fast GPU processing"],
    top_n=2,
    truncate="END",  # Nvidia-specific parameter
)

인증

Nvidia NIM API 키를 설정해요:

export NVIDIA_NIM_API_KEY="nvapi-..."

import os
os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."

# Or pass directly
response = litellm.rerank(
    model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
    query="test",
    documents=["doc1"],
    api_key="nvapi-...",
)

사용자 정의 API Base URL

기본 base URL을 여러 방식으로 덮어쓸 수 있어요:

옵션 1: 환경 변수

export NVIDIA_NIM_API_BASE="https://your-custom-endpoint.com"

옵션 2: 파라미터로 전달

response = litellm.rerank(
    model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
    query="test",
    documents=["doc1"],
    api_base="https://your-custom-endpoint.com",
)

옵션 3: 전체 URL (모델 경로 포함)

완전한 엔드포인트 URL이 있다면 직접 전달할 수 있어요:

response = litellm.rerank(
    model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
    query="test",
    documents=["doc1"],
    api_base="https://your-custom-endpoint.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking",
)

LiteLLM은 경로에 /retrieval/가 있는지 확인해 전체 URL을 감지하고 그대로 사용해요.

API 키는 어떻게 받나요?

Nvidia 웹사이트에서 Nvidia NIM API 키를 받을 수 있어요.

관련 문서

  • Nvidia NIM - 메인 문서
  • Nvidia NIM 채팅 완성
  • LiteLLM 리랭크 엔드포인트
  • Nvidia NIM 공식 문서 ↗

더 알아보기 (Learn more)

  • Nvidia NIM 리랭크 API 문서
  • LiteLLM 리랭크 기능