bge_m3_embedding_server/embedder/
sparse.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! BGE-M3 SPLADE-style sparse embedding pipeline.
16
17use anyhow::Result;
18use ort::value::TensorRef;
19
20use super::error::ort_err;
21use super::jit_guard::{self, TrtJitGuard};
22use super::math::{seq_len_distribution, sparse_maxpool, sparse_project};
23use super::tokenize::{build_chunk_arrays, tokenize_no_pad};
24use super::types::{EmbedStats, SparseEmbedding};
25use crate::binpack::{CostModel, bin_pack};
26use crate::config::ModelVariant;
27
28/// Produces sparse embeddings via the BGE-M3 sparse-linear projection.
29///
30/// Tokenizes once, then uses the cost model to bin-pack into chunks. Results
31/// are scattered back to the original input order.
32///
33/// `guard` is the in-band TRT JIT admission guard; see [`embed_dense`] for the
34/// refusal semantics.
35///
36/// [`embed_dense`]: super::dense::embed_dense
37#[allow(clippy::cast_possible_truncation)]
38pub(super) fn embed_sparse(
39    session: &mut ort::session::Session,
40    tokenizer: &tokenizers::Tokenizer,
41    texts: &[String],
42    cost_model: &CostModel,
43    model_variant: ModelVariant,
44    guard: Option<&TrtJitGuard>,
45) -> Result<(Vec<SparseEmbedding>, EmbedStats)> {
46    let (weight, bias) = crate::weights::sparse_linear();
47    let weight_view = weight.view();
48
49    let tokenize_start = std::time::Instant::now();
50    let encodings = tokenize_no_pad(tokenizer, texts)?;
51    let seq_lens: Vec<usize> = encodings.iter().map(|e| e.get_ids().len()).collect();
52    let tokenize_ms = u64::try_from(tokenize_start.elapsed().as_millis()).unwrap_or(u64::MAX);
53
54    let seq_dist = seq_len_distribution(&seq_lens);
55    let total_token_positions: usize = seq_lens.iter().sum();
56    let chunks = bin_pack(&seq_lens, cost_model);
57    jit_guard::guard_chunks(guard, &chunks, &seq_lens).map_err(anyhow::Error::new)?;
58
59    let mut all_sparse: Vec<Option<SparseEmbedding>> = (0..texts.len()).map(|_| None).collect();
60
61    let mut max_chunk_seq: usize = 0;
62    let mut inference_ms: u64 = 0;
63
64    for (chunk_idx, chunk_indices) in chunks.iter().enumerate() {
65        let chunk_max = chunk_indices
66            .iter()
67            .map(|&i| seq_lens[i])
68            .max()
69            .unwrap_or(1)
70            .max(1);
71
72        max_chunk_seq = max_chunk_seq.max(chunk_max);
73
74        let (ids_array, mask_array) = build_chunk_arrays(&encodings, chunk_indices, chunk_max)?;
75
76        let ids_tensor = TensorRef::from_array_view(ids_array.view()).map_err(ort_err)?;
77        let mask_tensor = TensorRef::from_array_view(mask_array.view()).map_err(ort_err)?;
78
79        let chunk_start = std::time::Instant::now();
80        let outputs = {
81            let _span = tracing::debug_span!(
82                "chunk",
83                chunk_idx,
84                batch = chunk_indices.len(),
85                max_seq = chunk_max
86            )
87            .entered();
88            session
89                .run(ort::inputs! {
90                    "input_ids" => ids_tensor,
91                    "attention_mask" => mask_tensor,
92                })
93                .map_err(ort_err)?
94        };
95        let chunk_ms = u64::try_from(chunk_start.elapsed().as_millis()).unwrap_or(u64::MAX);
96        inference_ms = inference_ms.saturating_add(chunk_ms);
97        tracing::debug!(
98            chunk_idx,
99            batch = chunk_indices.len(),
100            max_seq = chunk_max,
101            elapsed_ms = chunk_ms,
102            "sparse chunk inference complete"
103        );
104
105        // FP32: token_embeddings [batch, seq, 1024].
106        // FP16/INT8: last_hidden_state [batch, seq, 1024] — same shape, different key.
107        let token_emb = match model_variant {
108            ModelVariant::Fp32 => outputs["token_embeddings"]
109                .try_extract_array::<f32>()
110                .map_err(ort_err)?,
111            ModelVariant::Fp16 | ModelVariant::Int8 => outputs["last_hidden_state"]
112                .try_extract_array::<f32>()
113                .map_err(ort_err)?,
114        };
115
116        for (chunk_pos, &orig_idx) in chunk_indices.iter().enumerate() {
117            let enc = &encodings[orig_idx];
118            let ids = enc.get_ids();
119            let mask = enc.get_attention_mask();
120            let batch_hidden = token_emb.index_axis(ndarray::Axis(0), chunk_pos);
121
122            let scores: Vec<f32> = (0..ids.len())
123                .map(|j| {
124                    let hidden = batch_hidden.index_axis(ndarray::Axis(0), j);
125                    let hidden_slice = hidden
126                        .as_slice()
127                        .expect("hidden state should be contiguous");
128                    sparse_project(hidden_slice, &weight_view, *bias)
129                })
130                .collect();
131
132            let (indices, values) = sparse_maxpool(ids, mask, &scores);
133            all_sparse[orig_idx] = Some(SparseEmbedding { indices, values });
134        }
135    }
136
137    let stats = EmbedStats {
138        chunks: chunks.len(),
139        max_chunk_seq,
140        total_token_positions,
141        tokenize_ms,
142        inference_ms,
143        seq_len_min: seq_dist.min,
144        seq_len_max: seq_dist.max,
145        seq_len_mean: seq_dist.mean,
146        seq_len_p95: seq_dist.p95,
147    };
148
149    Ok((
150        all_sparse
151            .into_iter()
152            .map(|s| s.expect("every slot must be filled"))
153            .collect(),
154        stats,
155    ))
156}
bge_m3_embedding_server/embedder/sparse.rs

bge_m3_embedding_server/embedder/
sparse.rs
