# rag-reranker

**Repository Path**: wu-hen/rag-reranker

## Basic Information

- **Project Name**: rag-reranker
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-05-23
- **Last Updated**: 2025-05-23

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# vLLM 重排序服务

基于 FastAPI 与 vLLM 实现的文本重排序（Reranking）服务，支持通过 HTTP 接口对一条查询与多条文档进行相关性打分，并返回 Top N 的结果。

---

## 功能特性

* ✅ **跨编码器打分**：使用 `vLLM` 引擎对查询与候选文档进行交叉编码评分（`task='score'`）。
* 🔒 **简单鉴权**：支持 Bearer Token 认证，Token 值通过 `.env` 文件配置。
* ⚙️ **模型灵活**：支持从环境变量或请求体指定模型名称，兼容 Hugging Face Model Hub 或本地模型路径。
* 📦 **异步接口**：基于 FastAPI 异步设计，结合 `asyncio` & 线程池提升并发性能。
* 📑 **可扩展**：预留 Token 使用统计字段，可对接 vLLM 的实际 Token 统计信息。

---

## 快速开始

### 1. 克隆仓库

```bash
git clone https://your-repo-url.git
cd your-repo-folder
```

### 2. 安装依赖

```bash
pip install -r requirements.txt
```

### 3. 下载模型

如果使用 Hugging Face 上的模型，请在本地事先下载：

```python
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="BAAI/bge-reranker-v2-m3",
    local_dir="./models/bge-reranker-v2-m3"
)
```

下载后可在请求中使用本地路径 `./models/bge-reranker-v2-m3`。

### 4. 配置环境变量

项目根目录创建 `.env` 文件，示例：

```dotenv
MODEL_NAME=bge-reranker-v2-m3
AUTH_TOKEN=your_secure_token_here
```

* **MODEL\_NAME**：默认使用的模型，若在请求中未提供 `model` 字段，则使用该值。
* **AUTH\_TOKEN**：HTTP 接口的 Bearer Token，用于简单鉴权。

### 5. 启动服务

```bash
uvicorn main:app --host 0.0.0.0 --port 8000
```

服务将监听在 `http://0.0.0.0:8000`。

---

## 接口说明

### POST `/v1/rerank`

对单条查询与多条文档进行打分，返回 Top N 结果。

#### 请求头

```
Authorization: Bearer <AUTH_TOKEN>
Content-Type: application/json
```

#### 请求示例

```json
{
  "model": "bge-reranker-v2-m3",
  "query": "Apple",
  "documents": ["apple", "banana", "fruit", "vegetable"],
  "top_n": 4,
  "return_documents": false,
}
```

* `model`：可选，若未提供则使用 `.env` 中的 `MODEL_NAME`，也可填本地路径。
* `query`：查询文本。
* `documents`：候选文档列表。
* `top_n`：返回的结果条数。
* `return_documents`：是否在结果中包含原文档文本。

#### 响应示例

```json
{
  "id": "c6f8e4d2-3f2b-4a7e-8cb6-1d9e2f7b1234",
  "results": [
    {
      "document": { "text": "apple" },
      "index": 0,
      "relevance_score": 3.215
    },
    {
      "document": { "text": "fruit" },
      "index": 2,
      "relevance_score": 1.782
    }
  ],
  "tokens": { "input_tokens": 0, "output_tokens": 0 }
}
```

* `id`：请求的唯一 UUID。
* `results`：打分结果列表，按 `relevance_score` 降序排列。
* `tokens`：预留的 Token 使用统计。

---

## 测试

项目内提供异步 HTTP 测试用例 `test_rerank.py`，基于 `pytest` 和 `httpx`：

```bash
pip install pytest pytest-asyncio httpx
pytest test_rerank.py
```

---

## License

MIT