bge_m3_embedding_server/embedder/
dense.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//! Dense 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::{normalize_l2, seq_len_distribution};
23use super::tokenize::{build_chunk_arrays, tokenize_no_pad};
24use super::types::EmbedStats;
25use crate::binpack::{CostModel, bin_pack};
26use crate::config::ModelVariant;
27
28/// Produces L2-normalized dense embeddings.
29///
30/// Tokenizes once, then uses the cost model to bin-pack into chunks that fit
31/// within the workspace budget. Results are scattered back to the original
32/// input order.
33///
34/// `guard` is the in-band TRT JIT admission guard (see
35/// [`super::jit_guard`]). When `Some`, any chunk whose sequence length is in
36/// the dangerous range and uncovered by the warmed engine profile is refused
37/// (returning [`super::jit_guard::TrtJitRejection`]) before any `session.run()`
38/// executes, so a single request fails with HTTP `503` instead of risking a
39/// process-killing pathological autotuner allocation.
40#[allow(clippy::cast_possible_truncation)]
41pub(super) fn embed_dense(
42    session: &mut ort::session::Session,
43    tokenizer: &tokenizers::Tokenizer,
44    texts: &[String],
45    cost_model: &CostModel,
46    model_variant: ModelVariant,
47    guard: Option<&TrtJitGuard>,
48) -> Result<(Vec<Vec<f32>>, EmbedStats)> {
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    // Pre-allocate output slots (one per input text, filled below).
60    let mut all_embeddings: Vec<Vec<f32>> = (0..texts.len()).map(|_| Vec::new()).collect();
61
62    let mut max_chunk_seq: usize = 0;
63    let mut inference_ms: u64 = 0;
64
65    for (chunk_idx, chunk_indices) in chunks.iter().enumerate() {
66        let chunk_max = chunk_indices
67            .iter()
68            .map(|&i| seq_lens[i])
69            .max()
70            .unwrap_or(1)
71            .max(1); // guard: at least 1 to avoid 0-dim tensors
72
73        max_chunk_seq = max_chunk_seq.max(chunk_max);
74
75        let (ids_array, mask_array) = build_chunk_arrays(&encodings, chunk_indices, chunk_max)?;
76        let batch_len = ids_array.nrows();
77
78        let ids_tensor = TensorRef::from_array_view(ids_array.view()).map_err(ort_err)?;
79        let mask_tensor = TensorRef::from_array_view(mask_array.view()).map_err(ort_err)?;
80
81        let chunk_start = std::time::Instant::now();
82        let outputs = {
83            let _span = tracing::debug_span!(
84                "chunk",
85                chunk_idx,
86                batch = chunk_indices.len(),
87                max_seq = chunk_max
88            )
89            .entered();
90            session
91                .run(ort::inputs! {
92                    "input_ids" => ids_tensor,
93                    "attention_mask" => mask_tensor,
94                })
95                .map_err(ort_err)?
96        };
97        let chunk_ms = u64::try_from(chunk_start.elapsed().as_millis()).unwrap_or(u64::MAX);
98        inference_ms = inference_ms.saturating_add(chunk_ms);
99        tracing::debug!(
100            chunk_idx,
101            batch = chunk_indices.len(),
102            max_seq = chunk_max,
103            elapsed_ms = chunk_ms,
104            "dense chunk inference complete"
105        );
106
107        // FP32: sentence_embedding [batch, 1024] — pre-pooled CLS output.
108        // FP16/INT8: last_hidden_state [batch, seq, 1024] — CLS token at position 0.
109        let emb: ndarray::ArrayD<f32> = match model_variant {
110            ModelVariant::Fp32 => outputs["sentence_embedding"]
111                .try_extract_array::<f32>()
112                .map_err(ort_err)?
113                .to_owned(),
114            ModelVariant::Fp16 | ModelVariant::Int8 => {
115                let lhs = outputs["last_hidden_state"]
116                    .try_extract_array::<f32>()
117                    .map_err(ort_err)?;
118                lhs.index_axis(ndarray::Axis(1), 0).to_owned()
119            }
120        };
121
122        for (chunk_pos, &orig_idx) in chunk_indices.iter().enumerate() {
123            debug_assert!(chunk_pos < batch_len, "chunk_pos must be within batch");
124            let row = emb.index_axis(ndarray::Axis(0), chunk_pos);
125            let mut vec = row
126                .as_slice()
127                .expect("embedding should be contiguous")
128                .to_vec();
129            normalize_l2(&mut vec);
130            all_embeddings[orig_idx] = vec;
131        }
132    }
133
134    let stats = EmbedStats {
135        chunks: chunks.len(),
136        max_chunk_seq,
137        total_token_positions,
138        tokenize_ms,
139        inference_ms,
140        seq_len_min: seq_dist.min,
141        seq_len_max: seq_dist.max,
142        seq_len_mean: seq_dist.mean,
143        seq_len_p95: seq_dist.p95,
144    };
145
146    Ok((all_embeddings, stats))
147}
bge_m3_embedding_server/embedder/dense.rs

bge_m3_embedding_server/embedder/
dense.rs
