Livrare LOT 1 - Didi
This commit is contained in:
commit
5380c3fc63
990 changed files with 133308 additions and 0 deletions
274
ai_platform/modules/rerank/INDEX.md
Normal file
274
ai_platform/modules/rerank/INDEX.md
Normal file
|
|
@ -0,0 +1,274 @@
|
|||
# Rerank Service — INDEX
|
||||
|
||||
Cross-encoder reranking API for DIDI semantic search precision. Serves the `BAAI/bge-reranker-v2-m3` model behind a Cohere/Jina-compatible HTTP API. Used after kNN retrieval to score `(query, candidate)` pairs and reorder evidence by true relevance, instead of relying solely on pgvector cosine similarity from embeddings.
|
||||
|
||||
- **Stack:** Python 3.10+, FastAPI, Uvicorn, Pydantic v2, httpx, vLLM (or llama.cpp) backend
|
||||
- **URLs:**
|
||||
- Dev API: `http://10.11.10.12:54200`
|
||||
- Prod API: `http://10.11.10.12:14200`
|
||||
- vLLM internal: `:54201` (Dev) / `:14201` (Prod)
|
||||
- llama.cpp internal: `:54210` (Dev) / `:14210` (Prod)
|
||||
- **Containers:** `didiAI-rerank-api` (FastAPI wrapper) + `didiAI-rerank-vllm` (or `didiAI-rerank-llamacpp`)
|
||||
- **Model:** `BAAI/bge-reranker-v2-m3` (multilingual cross-encoder, ~8K context window)
|
||||
- **Port schema:** `x42xx` family (Reranking)
|
||||
|
||||
---
|
||||
|
||||
## Ce face
|
||||
|
||||
Re-scores candidate documents against a query using a cross-encoder model — directly attending over both texts together — for higher precision than embedding cosine similarity alone. The cross-encoder gives a true relevance score in `[0..1]`.
|
||||
|
||||
**Pipeline role inside DIDI:**
|
||||
|
||||
1. didi-brain `/v1/gather` runs an initial kNN retrieval over pgvector embeddings (dim 1024, separate `embeddings` module).
|
||||
2. The top-K candidates (typically 50–200) are passed to this rerank service.
|
||||
3. The reranker scores each `(query, candidate)` pair and returns a relevance-sorted list.
|
||||
4. didi-brain takes the top-N (e.g. top 10) reranked items as final evidence.
|
||||
|
||||
This two-stage retrieval (kNN → cross-encoder rerank) is the standard high-quality semantic search pattern: embeddings handle scale, the cross-encoder handles precision.
|
||||
|
||||
---
|
||||
|
||||
## API endpoints
|
||||
|
||||
Cohere/Jina-compatible. Auth is optional (Bearer token) — enabled if `RERANK_API_TOKENS` is set. Health endpoints stay public.
|
||||
|
||||
### `POST /v1/rerank` (alias `POST /v2/rerank`)
|
||||
|
||||
Rerank a list of documents against a query.
|
||||
|
||||
**Request body:**
|
||||
```json
|
||||
{
|
||||
"model": "BAAI/bge-reranker-v2-m3",
|
||||
"query": "What is machine learning?",
|
||||
"documents": ["Machine learning is...", "Cats are pets", "Deep learning..."],
|
||||
"top_n": 3,
|
||||
"return_documents": false,
|
||||
"backend": null
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-------|------|----------|-------------|
|
||||
| `model` | string | yes | Model identifier (e.g. `BAAI/bge-reranker-v2-m3`) |
|
||||
| `query` | string | yes | Search query |
|
||||
| `documents` | string[] | yes | Documents to rerank (1–1000) |
|
||||
| `top_n` | int | no | Return only top N results (default: all) |
|
||||
| `return_documents` | bool | no | Echo document text in response |
|
||||
| `backend` | string | no | Override default: `"vllm"` or `"llamacpp"` |
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"id": "rerank-abc123def456",
|
||||
"model": "BAAI/bge-reranker-v2-m3",
|
||||
"results": [
|
||||
{"index": 2, "relevance_score": 0.9523, "document": null},
|
||||
{"index": 0, "relevance_score": 0.8876, "document": null}
|
||||
],
|
||||
"usage": {"total_tokens": 150},
|
||||
"backend": "vllm"
|
||||
}
|
||||
```
|
||||
|
||||
### `GET /v1/models`
|
||||
|
||||
List loaded models. Optional `?backend=vllm|llamacpp` filter.
|
||||
|
||||
### `GET /v1/backends`
|
||||
|
||||
List enabled backends, e.g. `{"backends": ["vllm", "llamacpp"]}`.
|
||||
|
||||
### `GET /health`
|
||||
|
||||
Detailed per-backend health. Status is `healthy` / `degraded` / `unhealthy`.
|
||||
|
||||
### `GET /ready`
|
||||
|
||||
Simple K8s-style readiness probe — returns `{"ready": true}`.
|
||||
|
||||
### Error codes
|
||||
|
||||
`400` bad request · `401` auth failed · `429` rate-limited (with `Retry-After` header) · `503` backend unavailable · `504` backend timeout. All errors return `{"detail": "..."}`. Every response carries an `X-Request-ID` header for tracing.
|
||||
|
||||
Full reference in `API.md`.
|
||||
|
||||
---
|
||||
|
||||
## Backends
|
||||
|
||||
The API is a thin FastAPI router in front of one of two cross-encoder servers. Both backends speak OpenAI-compatible HTTP, so the wrapper unifies them.
|
||||
|
||||
| Backend | When | Pros | Cons |
|
||||
|---------|------|------|------|
|
||||
| **vLLM** | Production, GPU host | High throughput, batched scoring, lowest latency under load | Requires NVIDIA GPU + CUDA 12.x |
|
||||
| **llama.cpp** | Dev / CPU fallback | Runs on CPU or modest GPU, GGUF quantized models, low memory | Lower throughput |
|
||||
|
||||
The wrapper routes requests via a backend registry. Set `RERANK_DEFAULT_BACKEND` to choose, or override per-request with the `backend` field in the body. Both can be enabled simultaneously (`RERANK_ENABLE_VLLM=true`, `RERANK_ENABLE_LLAMACPP=true`).
|
||||
|
||||
vLLM is launched with `--task score` (vLLM 0.8.x cross-encoder mode). llama.cpp is launched with `--reranking`.
|
||||
|
||||
---
|
||||
|
||||
## How didi-brain uses it
|
||||
|
||||
didi-brain's gather/retrieval flow:
|
||||
|
||||
1. `services/gather.py` builds an initial candidate set via pgvector kNN over the `embeddings` module (1024-dim vectors).
|
||||
2. Calls the reranker via `shared/reranker_client.py` — typically wrapping the `/v1/rerank` endpoint with the brain's HTTPS/auth config.
|
||||
3. Picks the top-N reranked candidates as final evidence chunks.
|
||||
4. Passes them to the LLM as grounded context.
|
||||
|
||||
The reranker is therefore in the critical path of every brain `/v1/gather` call. Its latency budget is small (a few hundred ms), so vLLM batching matters in production.
|
||||
|
||||
---
|
||||
|
||||
## Structura fisiere
|
||||
|
||||
```
|
||||
rerank/
|
||||
├── README.md # Quick start, install, env vars table
|
||||
├── API.md # Full HTTP API reference + curl/Python examples
|
||||
├── INDEX.md # This file
|
||||
├── pyproject.toml # Package metadata, deps (fastapi, httpx, openai SDK)
|
||||
├── uv.lock # uv lockfile
|
||||
├── .env.example # All RERANK_* env vars documented
|
||||
├── deploy/
|
||||
│ ├── Dockerfile # API wrapper image
|
||||
│ ├── docker-compose.yml # api / vllm / llamacpp profiles
|
||||
│ └── deploy.sh # Helper: ./deploy.sh --profile vllm -d
|
||||
├── src/rerank/
|
||||
│ ├── __init__.py # Public exports (RerankClient, ...)
|
||||
│ ├── cli.py # `rerank` entrypoint — `python -m rerank.cli`
|
||||
│ ├── client.py # RerankClient (Python SDK to call this API)
|
||||
│ ├── config.py # Pydantic Settings, RERANK_ env prefix
|
||||
│ ├── schemas.py # Pydantic request/response models
|
||||
│ ├── exceptions.py # Custom error types
|
||||
│ ├── logging.py # Structured logging setup
|
||||
│ ├── types.py # Backend literal types
|
||||
│ ├── api/
|
||||
│ │ ├── app.py # FastAPI app factory
|
||||
│ │ ├── dependencies.py # Auth + rate-limit DI
|
||||
│ │ ├── middleware.py # X-Request-ID, rate limit, error handlers
|
||||
│ │ └── routes/
|
||||
│ │ ├── rerank.py # POST /v1/rerank, /v2/rerank
|
||||
│ │ ├── models.py # GET /v1/models, /v1/backends
|
||||
│ │ └── health.py # GET /health, /ready
|
||||
│ └── backends/
|
||||
│ ├── base.py # Abstract BackendBase (rerank, health, list_models)
|
||||
│ ├── vllm_backend.py # vLLM via OpenAI SDK (--task score)
|
||||
│ ├── llamacpp_backend.py # llama.cpp --reranking endpoint
|
||||
│ └── registry.py # Backend registry + selector
|
||||
└── tests/
|
||||
├── conftest.py
|
||||
├── test_config.py
|
||||
└── test_schemas.py
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Configuration
|
||||
|
||||
All env vars use the `RERANK_` prefix and are read via Pydantic Settings (`config.py`).
|
||||
|
||||
**Required (no defaults):**
|
||||
|
||||
| Var | Example |
|
||||
|-----|---------|
|
||||
| `RERANK_DEFAULT_BACKEND` | `vllm` or `llamacpp` |
|
||||
| `RERANK_ENABLE_VLLM` | `true` / `false` |
|
||||
| `RERANK_ENABLE_LLAMACPP` | `true` / `false` |
|
||||
| `RERANK_EXTERNAL_URL` | `http://localhost:14200` (used in OpenAPI spec) |
|
||||
|
||||
**API server:**
|
||||
|
||||
| Var | Default | Notes |
|
||||
|-----|---------|-------|
|
||||
| `RERANK_HOST` | `0.0.0.0` | |
|
||||
| `RERANK_PORT` | `14200` | Prod 14200 / Dev 54200 |
|
||||
| `RERANK_API_TOKENS` | (empty) | Comma-separated; if empty, auth is disabled |
|
||||
| `RERANK_RATE_LIMIT_RPS` | `20.0` | |
|
||||
| `RERANK_RATE_LIMIT_BURST` | `50` | |
|
||||
| `RERANK_MAX_CONCURRENT_RERANKS` | `20` | |
|
||||
| `RERANK_REQUEST_TIMEOUT` | `120.0` s | |
|
||||
| `RERANK_CONNECT_TIMEOUT` | `10.0` s | |
|
||||
| `RERANK_LOG_LEVEL` | `INFO` | |
|
||||
| `RERANK_LOG_JSON` | `false` | |
|
||||
|
||||
**vLLM backend:**
|
||||
|
||||
| Var | Default |
|
||||
|-----|---------|
|
||||
| `RERANK_VLLM_BASE_URL` | `http://localhost:54201` |
|
||||
| `RERANK_VLLM_MODEL` | `BAAI/bge-reranker-v2-m3` |
|
||||
| `RERANK_VLLM_PORT` | `14201` |
|
||||
| `RERANK_VLLM_GPU` | `0` |
|
||||
| `RERANK_VLLM_GPU_UTIL` | `0.50` |
|
||||
| `RERANK_VLLM_MAX_LEN` | `8192` |
|
||||
|
||||
**llama.cpp backend:**
|
||||
|
||||
| Var | Default |
|
||||
|-----|---------|
|
||||
| `RERANK_LLAMACPP_BASE_URL` | `http://localhost:54210` |
|
||||
| `RERANK_LLAMACPP_MODEL` | `bge-reranker-v2-m3-q4_k_m.gguf` |
|
||||
| `RERANK_LLAMACPP_PORT` | `14210` |
|
||||
| `RERANK_LLAMACPP_CTX` | `8192` |
|
||||
| `RERANK_LLAMACPP_THREADS` | `4` |
|
||||
| `RERANK_LLAMACPP_PARALLEL` | `4` |
|
||||
| `MODELS_DIR` | `/cai2_ds_storage/models` |
|
||||
|
||||
**Shared HF:** `HF_CACHE_DIR` (default `/cai2_ds_storage/hf_cache`), `HF_TOKEN`.
|
||||
|
||||
Full list with comments in `.env.example`.
|
||||
|
||||
---
|
||||
|
||||
## Performance / model
|
||||
|
||||
- **Model:** `BAAI/bge-reranker-v2-m3` — multilingual cross-encoder, ~8K context window, ~568M params.
|
||||
- **Output:** single relevance score per `(query, document)` pair, normalized to `[0..1]`.
|
||||
- **Expected latency:** roughly 100–300 ms for ~20 candidates per query on a single mid-range GPU with vLLM batching; scales sublinearly with batch size.
|
||||
- **Throughput:** dominated by GPU memory and `RERANK_VLLM_GPU_UTIL` / `RERANK_MAX_CONCURRENT_RERANKS`. vLLM continuous batching helps a lot under concurrent load.
|
||||
- **Memory:** vLLM defaults to 50% GPU memory util (`RERANK_VLLM_GPU_UTIL=0.50`), so it can co-host on a GPU shared with other modules. llama.cpp Q4_K_M quantization fits comfortably on CPU.
|
||||
|
||||
---
|
||||
|
||||
## Deployment
|
||||
|
||||
Docker Compose with three profiles in `deploy/docker-compose.yml`:
|
||||
|
||||
```bash
|
||||
cd /home/admin365/didi_mono/ai_platform/modules/rerank/deploy
|
||||
|
||||
# Production (GPU + vLLM)
|
||||
./deploy.sh --profile vllm -d
|
||||
|
||||
# Lightweight / no-GPU
|
||||
./deploy.sh --profile llamacpp -d
|
||||
|
||||
# API only (external rerank servers running elsewhere)
|
||||
./deploy.sh --profile api -d
|
||||
|
||||
# Logs / stop
|
||||
./deploy.sh --profile vllm --logs
|
||||
./deploy.sh --profile vllm --down
|
||||
```
|
||||
|
||||
Containers:
|
||||
- `didiAI-rerank-api` — FastAPI wrapper, port `RERANK_PORT` (14200/54200)
|
||||
- `didiAI-rerank-vllm` — `vllm/vllm-openai:v0.8.5`, GPU device pinned via `RERANK_VLLM_GPU`, internal port 14201
|
||||
- `didiAI-rerank-llamacpp` — `ghcr.io/ggml-org/llama.cpp:server-b4769`, internal port 8080 → host 14210/54210
|
||||
|
||||
All services share the external `didi-network` Docker network so other DIDI modules (didi-brain, embeddings) can reach them by container name.
|
||||
|
||||
Health gates: API has a 30-s `/health` healthcheck; vLLM has a 5-min start period to allow model load on first boot.
|
||||
|
||||
---
|
||||
|
||||
## Related
|
||||
|
||||
- **Embeddings module** (`ai_platform/modules/embeddings`) — separate service, dim 1024, used for the initial pgvector kNN stage that feeds this reranker.
|
||||
- **didi-brain** (`ai_platform/modules/didi_brain/brain_api`) — main consumer; `services/gather.py` calls rerank after kNN; `shared/reranker_client.py` is the helper.
|
||||
- **GPU host:** `10.11.10.17` / `10.11.10.12` (see project network architecture doc).
|
||||
Loading…
Add table
Add a link
Reference in a new issue