HuggingFace Rerank

HuggingFace Rerank

HuggingFace Rerank를 사용하면 Hugging Face 인프라 또는 자체 사용자 지정 엔드포인트에 호스팅된 reranking 모델로 쿼리에 대한 관련성을 기준으로 문서를 재정렬할 수 있어요.

출처: 문서

본문

개요 (Overview)

빠른 시작 (Quick Start)

LiteLLM Python SDK

import litellm
import os

# Set your HuggingFace token
os.environ["HUGGINGFACE_API_KEY"] = "hf_xxxxxx"

# Basic rerank usage
response = litellm.rerank(
    model="huggingface/BAAI/bge-reranker-base",
    query="What is the capital of the United States?",
    documents=[
        "Carson City is the capital city of the American state of Nevada.",
        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
        "Washington, D.C. is the capital of the United States.",
        "Capital punishment has existed in the United States since before it was a country.",
    ],
    top_n=3,
)
print(response)

사용자 지정 엔드포인트 사용

import litellm

response = litellm.rerank(
    model="huggingface/BAAI/bge-reranker-base",
    query="hello",
    documents=["hello", "world"],
    top_n=2,
    api_base="https://my-custom-hf-endpoint.com",
    api_key="test_api_key",
)
print(response)

비동기 사용법 (Async Usage)

import litellm
import asyncio
import os

os.environ["HUGGINGFACE_API_KEY"] = "hf_xxxxxx"

async def async_rerank_example():
    response = await litellm.arerank(
        model="huggingface/BAAI/bge-reranker-base",
        query="What is the capital of the United States?",
        documents=[
            "Carson City is the capital city of the American state of Nevada.",
            "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
            "Washington, D.C. is the capital of the United States.",
            "Capital punishment has existed in the United States since before it was a country.",
        ],
        top_n=3,
    )
    print(response)

asyncio.run(async_rerank_example())

LiteLLM Proxy

1. config.yaml에서 모델 설정

model_list:
  - model_name: bge-reranker-base
    litellm_params:
      model: huggingface/BAAI/bge-reranker-base
      api_key: os.environ/HF_TOKEN
  - model_name: bge-reranker-large
    litellm_params:
      model: huggingface/BAAI/bge-reranker-large
      api_key: os.environ/HF_TOKEN
  - model_name: custom-reranker
    litellm_params:
      model: huggingface/BAAI/bge-reranker-base
      api_base: https://my-custom-hf-endpoint.com
      api_key: your-custom-api-key

2. Proxy 시작

export HF_TOKEN="hf_xxxxxx"
litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000

3. rerank 요청

curl:

curl http://localhost:4000/rerank \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer ***" \
  -d '{
    "model": "bge-reranker-base",
    "query": "What is the capital of the United States?",
    "documents": [
      "Carson City is the capital city of the American state of Nevada.",
      "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
      "Washington, D.C. is the capital of the United States.",
      "Capital punishment has existed in the United States since before it was a country."
    ],
    "top_n": 3
  }'

Python SDK:

import litellm

# Initialize with your LiteLLM proxy URL
response = litellm.rerank(
    model="bge-reranker-base",
    query="What is the capital of the United States?",
    documents=[
        "Carson City is the capital city of the American state of Nevada.",
        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
        "Washington, D.C. is the capital of the United States.",
        "Capital punishment has existed in the United States since before it was a country.",
    ],
    top_n=3,
    api_base="http://localhost:4000",
    api_key="your-litellm-api-key"
)
print(response)

requests 라이브러리:

import requests

url = "http://localhost:4000/rerank"
headers = {
    "Authorization": "Bearer your-litellm-api-key",
    "Content-Type": "application/json"
}
data = {
    "model": "bge-reranker-base",
    "query": "What is the capital of the United States?",
    "documents": [
        "Carson City is the capital city of the American state of Nevada.",
        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
        "Washington, D.C. is the capital of the United States.",
        "Capital punishment has existed in the United States since before it was a country."
    ],
    "top_n": 3
}
response = requests.post(url, headers=headers, json=data)
print(response.json())

설정 옵션 (Configuration Options)

인증 (Authentication)

HuggingFace 토큰(Serverless) 사용:

import os

os.environ["HUGGINGFACE_API_KEY"] = "hf_xxxxxx"

# Or pass directly
litellm.rerank(
    model="huggingface/BAAI/bge-reranker-base",
    api_key="hf_xxxxxx",
    # ... other params
)

사용자 지정 엔드포인트 사용:

litellm.rerank(
    model="huggingface/BAAI/bge-reranker-base",
    api_base="https://your-custom-endpoint.com",
    api_key="your-custom-key",
    # ... other params
)

응답 형식:

{
  "results": [
    { "index": 3, "relevance_score": 0.999071 },
    { "index": 4, "relevance_score": 0.7867867 },
    { "index": 0, "relevance_score": 0.32713068 }
  ],
  "id": "07734bd2-2473-4f07-94e1-0d9f0e6843cf",
  "meta": {
    "api_version": {
      "version": "2",
      "is_experimental": false
    },
    "billed_units": {
      "search_units": 1
    }
  }
}

더 알아보기 (Learn more)

• LiteLLM 개요
• LiteLLM Python SDK 사용법
• LiteLLM Proxy Server 설정
• 지원되는 공급업체 전체 목록