Metadata-Version: 2.4
Name: sg-reranker
Version: 0.1.0
Summary: Lightweight client for the SG reranker service
License: Proprietary
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: requests>=2.28
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"

# BGE Reranker v2-m3 — Docker Service

A production-ready CPU reranker microservice running
[BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3)
via ONNX Runtime (INT8 dynamic quantization).

Tuned for a GCP **n2-highmem-2** VM: 2 vCPUs, 16 GB RAM.

---

## Quick start

### 1 — Build

```bash
docker compose build
```

The build downloads the model from Hugging Face and bakes the quantized ONNX
file into the image (≈ 5–15 min, depending on network speed and CPU).
Subsequent rebuilds are fast because Docker caches the conversion layer as long
as `requirements.txt` and `app/convert.py` are unchanged.

### 2 — Configure (optional)

```bash
cp .env.example .env
# Edit .env to set API_KEY, MAX_LENGTH, thread counts, etc.
```

### 3 — Run

```bash
docker compose up -d
docker compose logs -f   # watch startup + warmup
```

The service is ready when you see **"Model ready — serving on port 8000"** in
the logs (usually 10–30 s after container start).

### 4 — Health check

```bash
curl http://localhost:8000/health
# {"status":"ok","model":"BAAI/bge-reranker-v2-m3","max_length":512}
```

### 5 — Rerank

```bash
curl -s -X POST http://localhost:8000/rerank \
  -H "Content-Type: application/json" \
  -d '{
    "query": "What is machine learning?",
    "documents": [
      "Machine learning enables systems to learn from data without explicit programming.",
      "Python is a popular general-purpose programming language.",
      "Deep learning models complex patterns using multi-layer neural networks."
    ],
    "top_n": 2,
    "return_documents": true
  }' | python3 -m json.tool
```

With bearer-token auth enabled (`API_KEY` set in `.env`):

```bash
curl -s -X POST http://localhost:8000/rerank \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $API_KEY" \
  -d '{...}'
```

### 6 — Python client

```bash
# against local container
python client.py

# against a remote VM
RERANKER_URL=http://10.0.0.5:8000 API_KEY=secret python client.py
```

---

## API reference

### `GET /health`

```json
{"status": "ok", "model": "BAAI/bge-reranker-v2-m3", "max_length": 512}
```

### `POST /rerank`

**Request body**

| Field              | Type           | Default | Description                              |
|--------------------|----------------|---------|------------------------------------------|
| `query`            | string         | —       | The search query                         |
| `documents`        | list of string | —       | Candidate documents to score             |
| `top_n`            | int \| null    | null    | Return only the top N results            |
| `max_length`       | int \| null    | null    | Token limit per pair (env var fallback)  |
| `return_documents` | bool           | true    | Include document text in the response    |

**Response**

```json
{
  "results": [
    {"index": 0, "score": 0.9821, "document": "Machine learning enables …"},
    {"index": 2, "score": 0.7634, "document": "Deep learning models …"}
  ],
  "model": "BAAI/bge-reranker-v2-m3"
}
```

Results are sorted by `score` descending. `index` refers to the original
position in the `documents` list so you can map back to your data.

---

## Reaching the service from another machine

There are three common approaches; choose based on your threat model.

### Option A — Same-VPC internal IP (recommended for GCP)

If your client runs in the same GCP VPC as the VM, use the VM's **internal
IP** directly. No firewall rule is needed because traffic stays on the private
network.

```bash
RERANKER_URL=http://10.128.0.X:8000 python client.py
```

Find the internal IP: `gcloud compute instances describe <VM_NAME> --format='get(networkInterfaces[0].networkIP)'`

### Option B — External IP with a locked-down firewall rule

Create a GCP firewall rule that allows TCP port 8000 only from specific source
IP ranges (your office CIDR, a bastion IP, etc.):

```bash
gcloud compute firewall-rules create allow-reranker \
  --direction=INGRESS \
  --action=ALLOW \
  --rules=tcp:8000 \
  --source-ranges=<YOUR_CIDR> \
  --target-tags=reranker-vm
```

Then tag your VM and hit it on its external IP:

```bash
RERANKER_URL=http://<EXTERNAL_IP>:8000 python client.py
```

> **Warning** — plain HTTP is unencrypted.  Anyone on the network path can
> read requests and responses, including the documents you are reranking.
> Use this option only within a trusted network or with TLS termination
> (e.g. a load balancer or nginx with a certificate).

### Option C — SSH tunnel (for local development)

Forward a local port to the service through an encrypted SSH session.  No
firewall rule is needed and no traffic leaves the tunnel unencrypted.

```bash
# In one terminal — keep this open
gcloud compute ssh <VM_NAME> -- -N -L 8000:localhost:8000

# In another terminal
curl http://localhost:8000/health
python client.py          # uses http://localhost:8000 by default
```

---

## Environment variables

| Variable          | Default                         | Description                              |
|-------------------|---------------------------------|------------------------------------------|
| `MODEL_DIR`       | `/models/onnx_reranker_quant`   | Path to the quantized ONNX model dir     |
| `MAX_LENGTH`      | `512`                           | Default token limit per query+doc pair   |
| `OMP_NUM_THREADS` | `2`                             | PyTorch / OpenBLAS thread count          |
| `ORT_NUM_THREADS` | `2`                             | ONNX Runtime intra-op thread count       |
| `API_KEY`         | _(unset = auth disabled)_       | Bearer token for `/rerank`               |

---

## Notes

- **Single Uvicorn worker** is intentional.  Each worker loads a full copy of
  the model.  On a 2-vCPU VM a second worker would double RAM usage and cause
  thread contention rather than improve throughput.
- **INT8 quantization** cuts model size by ~3× and speeds up matrix
  multiplications on CPUs that support AVX2 or AVX-512 VNNI (most modern Intel
  and AMD cores).  Accuracy loss on typical reranking benchmarks is < 1%.
- **TLS** — this service speaks plain HTTP.  For production deployments on the
  public internet, terminate TLS at a load balancer or a reverse proxy (nginx,
  Caddy) and keep the container on an internal network.
