embeddings: native bge-m3 dense+sparse(+colbert) multi-vector support

Add BAAI/bge-m3 hybrid embedding (dense + sparse lexical, optional colbert)
without the FlagEmbedding dependency (which pins older transformers/peft and
would conflict with the transformers 5.x stack): the sparse/colbert heads are
plain Linear layers over XLM-RoBERTa's last hidden state, reproduced natively
with transformers + torch already in the venv.

- new 'bge-m3' embedding backend: detected by xlm-roberta model_type + a
  sparse_linear.pt head (so plain bge-large etc. stay on the dense ST path);
  loads the base model + sparse_linear.pt (+ colbert_linear.pt if present).
- _bge_m3_encode: dense = L2-normalized CLS; sparse = max-pooled
  relu(sparse_linear·h) per token id (specials dropped) as {token_id: weight};
  colbert = L2-normalized colbert_linear·h per content token.
- /v1/embeddings extended: request.embedding_types selects any of
  dense|sparse|colbert (overrides the model-config default); response items
  gain sparse_embedding {indices, values} and colbert_embedding. Dense stays
  in the standard `embedding` field (base64 honored). Configurable default via
  the model's embedding_types config.
- cleanup()/eviction branch moves the base model + heads off GPU.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Mw2KQiswmD69T45fTfjKwW
parent 2c419247
This diff is collapsed.
......@@ -27,6 +27,7 @@ class EmbeddingsRequest(BaseModel):
encoding_format: Optional[str] = Field("float", description="Return embeddings as 'float' arrays or 'base64'.")
dimensions: Optional[int] = Field(None, description="Truncate embeddings to N dimensions (if the model supports it).")
quantization: Optional[str] = Field(None, description="Optional TurboQuant vector quantization: 'turbo' (8-bit), 'turbo8', 'turbo6', 'turbo4' or 'turbo2'. With encoding_format='float' the (lossy) reconstructed vectors are returned; with 'base64' the compact packed bytes are returned plus a 'quantization' metadata block describing how to decode them.")
embedding_types: Optional[List[str]] = Field(None, description="For multi-vector models (e.g. BAAI/bge-m3): which vector types to compute — any of 'dense', 'sparse', 'colbert'. Overrides the model's configured default. 'dense' → the usual `embedding`; 'sparse' → `sparse_embedding` {indices, values}; 'colbert' → `colbert_embedding` (token-level multi-vector).")
user: Optional[str] = Field(None, description="Opaque end-user identifier (passthrough).")
model_config = ConfigDict(extra="allow")
......@@ -34,7 +35,10 @@ class EmbeddingsRequest(BaseModel):
class EmbeddingObject(BaseModel):
object: str = "embedding"
index: int
embedding: Union[List[float], str] # float list or base64
embedding: Union[List[float], str] # float list or base64 (dense; [] if dense not requested)
# Multi-vector outputs (bge-m3 and similar). Present only when requested.
sparse_embedding: Optional[Dict] = None # {"indices": [token_id...], "values": [float...]}
colbert_embedding: Optional[List[List[float]]] = None # token-level vectors
class EmbeddingsResponse(BaseModel):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment