bge_m3_embedding_server/probe/
corpus.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//! Probe text synthesis helpers.
16//!
17//! The probe sweeps `(batch, seq)` shapes by submitting synthesized texts to
18//! the leader worker. Texts come from the curated benchmark corpus; we
19//! repeat/trim corpus entries to hit the target token count for each shape.
20
21/// Loads the benchmark corpus for use as probe text material.
22///
23/// Falls back to a tiny built-in sentence if the corpus file is not found.
24pub(super) fn load_probe_texts() -> Vec<String> {
25    let corpus_path =
26        std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("benches/fixtures/corpus.json");
27    if let Ok(raw) = std::fs::read_to_string(&corpus_path)
28        && let Ok(json) = serde_json::from_str::<serde_json::Value>(&raw)
29        && let Some(scenarios) = json["scenarios"].as_object()
30    {
31        let mut texts: Vec<String> = Vec::new();
32        for scenario in scenarios.values() {
33            if let Some(arr) = scenario["texts"].as_array() {
34                texts.extend(arr.iter().filter_map(|v| v.as_str().map(String::from)));
35            }
36        }
37        if !texts.is_empty() {
38            return texts;
39        }
40    }
41    // Fallback: minimal probe text.
42    vec![
43        "The embedding server startup probe synthesizes texts to measure workspace cost."
44            .to_string(),
45    ]
46}
47
48/// Synthesizes `batch` texts each of approximately `target_seq` tokens.
49///
50/// Token estimation: ~4 chars/token for natural English text.
51/// We repeat/trim corpus texts to hit the target character count.
52pub(super) fn synthesize_texts(corpus: &[String], batch: usize, target_seq: usize) -> Vec<String> {
53    let target_chars = target_seq.saturating_mul(4).max(16);
54    (0..batch)
55        .map(|i| {
56            let base = &corpus[i % corpus.len()];
57            // Repeat the base text until we have enough characters.
58            let repeated = base.repeat((target_chars / base.len().max(1)).max(2) + 1);
59            // Trim to target_chars bytes (not chars, but close enough for probing).
60            let trimmed = if repeated.len() > target_chars {
61                &repeated[..target_chars]
62            } else {
63                &repeated
64            };
65            trimmed.to_string()
66        })
67        .collect()
68}
bge_m3_embedding_server/probe/corpus.rs

bge_m3_embedding_server/probe/
corpus.rs
