feat(server): late-interaction scoring endpoint for multi-vector embedders
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
PR #1414 added the late-interaction (ColBERT MaxSim) score, `crate::embeddings::maxsim` and `maxsim_mlx` (`src/embeddings/maxsim.rs:42` and `:74`, re-exported at `src/embeddings/mod.rs:73`), and `mlxcel embed` prints a MaxSim similarity matrix for ColIdefics3 and ColQwen2.5 (`src/commands/embed.rs:89-100`, `similarity`). Nothing over HTTP consumes it. `POST /v1/embeddings` returns the raw `[num_tokens, 128]` matrices (for the PR's two-page example: `[[24, 128], [876, 128], [876, 128]]` for ColIdefics3, 1776 rows in total), so a client that wants "which page answers this query" has to download every document's token matrix and run MaxSim itself. The epic marked a scoring endpoint out of scope; this issue adds it.
## Current Behavior
- Reranking is served by `POST /v1/rerank` and the `/rerank` alias (`src/server/app.rs:126` and `:160`) through `create_rerank` (`src/server/routes/rerank.rs:141-200`), which validates a `RerankRequest` (`src/server/types/rerank.rs:27-48`: `model`, `query: RerankInput`, `documents: Vec`, `top_n`, `return_documents`, `instruction`), converts items with `to_rerank_item` (text or fetched image, `rerank.rs:104-139`), and calls `RerankModelProvider::rerank(query, documents, instruction)` (`src/server/rerank_model.rs:53-72`) inside `spawn_blocking`. The response is `RerankResponse { model, results: Vec, usage: RerankUsage { prompt_tokens, total_tokens } }` (`types/rerank.rs:129-152`). With no reranker loaded the route returns 501 `NO_RERANKER_MODEL_MESSAGE` (`rerank.rs:44` and `:147-149`).
- The reranker kinds are `RerankerKind::{SequenceClassifier, GenerativeText, GenerativeVl}` (`src/rerank/mod.rs:74-81`); `RerankItem { text: Option, image: Option }` (`:119-125`); `RerankScores { scores: Vec, prompt_tokens }` (`:167-172`). The provider is `RerankWorkerProvider` (`src/server/rerank_worker.rs:285`, `load` at `:305`), which owns its own worker thread and weights.
- Embeddings are served by `EmbeddingModelProvider` (`src/server/embedding_model.rs:64-100`: `embed_texts(Vec, EmbedOptions)`, `embed_tokens`, `embed_image(ImageInput, EmbedOptions)`, `model_id`, `multi_vector() -> bool`, `supports_images()`), returning `EmbedReply { vectors: Vec, prompt_tokens }` (`src/embeddings/engine.rs:89-93`) where `EmbeddingVector { values, shape }` has `is_multi_vector()` and `rows()` (`engine.rs:63-80`).
- `AppState` holds `embedding_model: Option>` and `rerank_model: Option>` (`src/server/state.rs:416` and `:421`).
## Proposed Solution
Extend `/v1/rerank` rather than add a new route, so clients keep one request and response schema for "score documents against a query".
1. Add `MaxSimRerankProvider` in a new `src/server/rerank_maxsim.rs`: `pub struct MaxSimRerankProvider { embedding: Arc }` implementing `RerankModelProvider`. `rerank(query, documents, instruction)` embeds the query as text (or image, if the query item carries one and the embedder `supports_images()`), embeds each document with `embed_texts` or `embed_image`, runs `crate::embeddings::maxsim(query_rows, doc_rows)` per document, and returns `RerankScores { scores, prompt_tokens }` with `prompt_tokens` equal to the sum of the replies' `prompt_tokens`. `kind()` returns a new `RerankerKind::LateInteraction` variant added to `src/rerank/mod.rs:74-81` (`accepts_instruction()` returns `false` for it; `supports_images()` mirrors the embedder). `model_id()` and `created_at()` come from the embedding provider. The `instruction` argument is rejected the same way the route already rejects it for sequence classifiers (`rerank.rs:172-179`), so the route needs no new branch.
2. Wiring in `src/server/startup.rs`: after the embedding provider is built, if `embedding.multi_vector()` is `true` and no `--reranker-model` was given, install `Arc::new(MaxSimRerankProvider { embedding })` through `AppState::with_rerank_model`. If `--reranker-model` was given as well, the explicit reranker wins and the startup log says the multi-vector embedder is not exposed on `/v1/rerank`. Log the installed provider the way the other side models are logged.
3. Scores are raw MaxSim sums (scale is the query length, as `src/commands/embed.rs:83-88` documents), not probabilities. Document this in the `RerankResult.relevance_score` doc comment in `types/rerank.rs` and in `docs/` where `/v1/rerank` is described; do not normalize, because the ordering and the CLI parity are what clients need and a normalization would hide the margin.
4. `/v1/models` lists the served model id once (the embedder), not a second entry for the MaxSim provider; check where the rerank id is added to the models list and skip it when the id equals the embedder's.
Rejected alternative: a separate `POST /v1/score` route. It would duplicate `RerankRequest` validation, image fetching and the response envelope for no client benefit; the `kind()` mechanism already lets one route serve three reranker kinds.
## Scope
**In scope:** `src/server/rerank_maxsim.rs` (new) and its tests, `src/rerank/mod.rs` (`RerankerKind::LateInteraction`), `src/server/startup.rs`, `src/server/routes/rerank.rs` only if a kind-specific message is needed, `src/server/types/rerank.rs` doc comment, `docs/` for `/v1/rerank`, a real-checkpoint gate.
**Out of scope:** batching documents through the embedding worker in one call (each document is one `embed_*` call; measure first), MaxSim on the GPU via `maxsim_mlx` inside the worker (the read-back rows are already on the route side, and the CPU `maxsim` is what the CLI uses), a new response schema.
## Implementation Notes
- **Reuse**: `crate::embeddings::maxsim` for the score; `RerankModelProvider` and `RerankScores` so `create_rerank` needs no change to its flow; `to_rerank_item` and `validate_items` for image fetching and shape checks; the route test harness `app_with` in `src/server/routes/rerank_tests.rs:54-66` with a stub `EmbeddingModelProvider`; `local_checkpoint` (`src/models/embedding_test_support.rs:231`) and `mlx_test_guard` for the real-checkpoint gate.
- **Constraints**: the provider must never be installed for a single-vector embedder (`multi_vector() == false`), otherwise MaxSim degenerates to a dot product that is not what `/v1/rerank` promises; `top_n` and `return_documents` continue to be applied by the route after scoring; the call is blocking and must stay inside the route's `spawn_blocking`.
- **Edge cases**: an image document on a text-only multi-vector embedder (400 through `validate_items`, as today); an empty document list (already 400); a document whose embedding has zero rows (score 0.0, matching `similarity` in `embed.rs:91-94`); query and documents with different token counts (MaxSim is asymmetric by design; the query is always the outer sum).
- **Error handling**: `EmbeddingError` from the embedder maps to `RerankError::Internal(msg)` (`src/server/rerank_model.rs:44-49`) so the route returns the existing 500 envelope; queue-full maps to `RerankError::QueueFull` and the existing 503.
## Acceptance Criteria
- [ ] `mlxcel-server -m ` installs the MaxSim provider and `POST /v1/rerank` with `{"query":"...","documents":[{"image_url":"data:image/png;base64,..."},{"image_url":"..."}]}` returns `RerankResponse` with one `relevance_score` per document, in document order, sorted by the route's existing `top_n` handling.
- [ ] A route test with a stub multi-vector `EmbeddingModelProvider` (two documents with known rows) asserts the returned `relevance_score`s equal `crate::embeddings::maxsim` of the same rows and `usage.prompt_tokens` equals the summed reply tokens.
- [ ] A route test asserts that `instruction` on the MaxSim provider returns 400 with the existing message, and that a single-vector stub embedder never gets the provider installed.
- [ ] A real-checkpoint gate on the merged colSmol-256M checkpoint (`vidore/ColSmolVLM-Instruct-256M-base` plus the `vidore/colSmol-256M` adapter merged, the same directory `src/models/colidefics3_tests.rs:567` uses) scores the PR #1414 query against its matching and unrelated pages and asserts `18.7396` against `8.6879` within 1e-2, the CLI's ordering, soft-skipping when the directory is absent.
- [ ] Integrated into the real code flow: the provider is installed from `startup.rs`, not only constructed in tests, and `docs/` for `/v1/rerank` describes the late-interaction kind and the raw-score scale.
## Verification
```bash
cargo fmt --all -- --check
cargo clippy --lib --tests --profile test-fast --features cuda -- -D warnings
cargo test --profile test-fast --features cuda --lib server::rerank_maxsim
cargo test --profile test-fast --features cuda --lib server::routes::rerank
cargo test --profile test-fast --features cuda --lib -- --test-threads=1 --ignored colsmol_maxsim_rerank_route_reproduces_cli_ordering
```
Manual: `mlxcel-server -m ` then `curl -s localhost:8080/v1/rerank -H 'content-type: application/json' -d @request.json | jq '.results[].relevance_score'`; pass is the matching page scoring about 18.74 and the unrelated page about 8.69.
## Technical Considerations
Related: PR #1414 (MaxSim, ColIdefics3, ColQwen2.5, and the measured 18.7396 vs 8.6879 numbers), PR #1417 (`/v1/rerank` and the reranker provider it extends), `src/models/col_late_interaction.rs` (adapter-merge handling for the col* checkpoints).
Originating PR: #1414.
Follow-up from epic #1348.
Contributor guide
Research direction
Start with src/server/rerank_model.rs, src/server/embedding_model.rs, src/server/routes/rerank.rs, and startup wiring; compare the existing provider implementations and rerank_tests.rs harness. Add the MaxSim provider and LateInteraction kind, then cover startup selection, scores, prompt-token usage, instruction rejection, model listing, and documentation. Run the listed rerank tests and the optional real-checkpoint test to verify CLI ordering and expected scores.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- api, backend, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100