diff --git a/colgrep/src/index/mod.rs b/colgrep/src/index/mod.rs index 559b5fe6..0144bdb7 100644 --- a/colgrep/src/index/mod.rs +++ b/colgrep/src/index/mod.rs @@ -25,7 +25,9 @@ use serde::{Deserialize, Serialize}; use crate::acceleration::apply_acceleration_mode; use crate::acceleration::{env_acceleration_mode_lossy, AccelerationMode}; use crate::embed::build_embedding_text; -use crate::parser::{build_call_graph, detect_language, extract_units, CodeUnit, Language}; +use crate::parser::{ + build_call_graph, detect_language, extract_units, CodeUnit, Language, UnitType, +}; use crate::signal::{is_interrupted, is_interrupted_outside_critical, CriticalSectionGuard}; use paths::{ @@ -56,6 +58,90 @@ const BUILDING_MARKER: &str = ".building"; /// more often (more resumable) at the cost of more index-append overhead. const BUILD_CHECKPOINT_UNITS: usize = 4096; +/// Upper bound for adaptive indexing with the default code model. Most code units use much +/// less based on their AST-derived shape; complex functions can use the full budget. +const ADAPTIVE_INDEX_DOCUMENT_LENGTH: usize = 512; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum IndexDocumentLength { + Native, + Fixed(usize), + Adaptive, +} + +fn build_checkpoint_units(binary: bool, index_exists: bool) -> usize { + if binary || !index_exists { + usize::MAX + } else { + BUILD_CHECKPOINT_UNITS + } +} + +fn index_document_length(model_id: &str) -> IndexDocumentLength { + let override_value = std::env::var("COLGREP_INDEX_DOCUMENT_LENGTH").ok(); + resolve_index_document_length(model_id, override_value.as_deref()) +} + +fn resolve_index_document_length( + model_id: &str, + override_value: Option<&str>, +) -> IndexDocumentLength { + if let Some(value) = override_value { + if value == "0" { + return IndexDocumentLength::Native; + } + if let Ok(length) = value.parse::() { + return IndexDocumentLength::Fixed(length.max(2)); + } + } + if model_id == crate::model::DEFAULT_MODEL { + IndexDocumentLength::Adaptive + } else { + IndexDocumentLength::Native + } +} + +fn model_document_length(policy: IndexDocumentLength) -> Option { + match policy { + IndexDocumentLength::Native => None, + IndexDocumentLength::Fixed(length) => Some(length), + IndexDocumentLength::Adaptive => Some(ADAPTIVE_INDEX_DOCUMENT_LENGTH), + } +} + +/// Choose how much source context is worth embedding for one AST-derived code unit. +/// +/// Identity and signature metadata come first in `build_embedding_text`, so simple symbols +/// need little context. Branching and high-complexity functions retain progressively more of +/// their bodies. Non-AST documents use their line span as the only available structure signal. +fn code_unit_document_length(unit: &CodeUnit, policy: IndexDocumentLength) -> usize { + match policy { + IndexDocumentLength::Native => usize::MAX, + IndexDocumentLength::Fixed(length) => length, + IndexDocumentLength::Adaptive => match unit.unit_type { + UnitType::Constant | UnitType::Class => 64, + UnitType::Function | UnitType::Method => { + if unit.complexity >= 10 { + 512 + } else if unit.complexity >= 5 || unit.has_error_handling { + 256 + } else if unit.complexity >= 2 || unit.has_loops || unit.has_branches { + 128 + } else { + 64 + } + } + UnitType::RawCode => match unit.end_line.saturating_sub(unit.line) { + 0..=15 => 64, + 16..=63 => 128, + _ => 256, + }, + UnitType::Section => 128, + UnitType::Document => 256, + }, + } +} + /// Test-only counter of expensive `delete_from_index` invocations. /// /// Issue #116: deleting many files used to call the full-index-rewrite primitive once per @@ -403,6 +489,7 @@ pub struct UpdatePlan { struct SortedUnit { unit: Arc, text: Arc, + max_length: usize, } /// After deduplication: unique texts to encode + a map from original positions @@ -410,6 +497,7 @@ struct SortedUnit { struct PreparedChunk { units: Vec>, unique_texts: Vec>, + unique_max_lengths: Vec, original_to_unique: Vec, } @@ -457,12 +545,17 @@ struct ChunkForCoding { /// When encoding more than this many units, prompt the user unless auto_confirm is set. pub const CONFIRMATION_THRESHOLD: usize = 30_000; -fn prepare_units_for_encoding(units: &[CodeUnit], sample_prefix_size: usize) -> Vec { +fn prepare_units_for_encoding( + units: &[CodeUnit], + sample_prefix_size: usize, + document_length: IndexDocumentLength, +) -> Vec { let mut items: Vec = units .iter() .map(|unit| SortedUnit { unit: Arc::new(unit.clone()), text: Arc::::from(build_embedding_text(unit)), + max_length: code_unit_document_length(unit, document_length), }) .collect(); @@ -500,17 +593,20 @@ fn prepare_units_for_encoding(units: &[CodeUnit], sample_prefix_size: usize) -> /// unit can retrieve its embedding after the GPU pass. On large codebases /// this saves ~8% of encoding work (e.g. re-exported types, trait impls). fn prepare_deduplicated_chunk(unit_chunk: &[SortedUnit]) -> PreparedChunk { - let mut index_by_text: HashMap<&str, usize> = HashMap::new(); + let mut index_by_text: HashMap<(&str, usize), usize> = HashMap::new(); let mut unique_texts: Vec> = Vec::new(); + let mut unique_max_lengths: Vec = Vec::new(); let mut original_to_unique: Vec = Vec::with_capacity(unit_chunk.len()); for item in unit_chunk.iter() { - if let Some(&unique_idx) = index_by_text.get(item.text.as_ref()) { + let key = (item.text.as_ref(), item.max_length); + if let Some(&unique_idx) = index_by_text.get(&key) { original_to_unique.push(unique_idx); } else { let unique_idx = unique_texts.len(); - index_by_text.insert(item.text.as_ref(), unique_idx); + index_by_text.insert(key, unique_idx); unique_texts.push(Arc::clone(&item.text)); + unique_max_lengths.push(item.max_length); original_to_unique.push(unique_idx); } } @@ -521,6 +617,7 @@ fn prepare_deduplicated_chunk(unit_chunk: &[SortedUnit]) -> PreparedChunk { .map(|item| Arc::clone(&item.unit)) .collect(), unique_texts, + unique_max_lengths, original_to_unique, } } @@ -574,7 +671,8 @@ fn run_tokenize_stage( .iter() .map(|text| text.as_ref()) .collect(); - let prepared_batches = model.tokenize_documents_in_batches(&text_refs)?; + let prepared_batches = model + .tokenize_documents_in_batches_with_limits(&text_refs, &chunk.unique_max_lengths)?; sender .send(TokenizedChunk { @@ -1327,6 +1425,7 @@ impl IndexBuilder { let batch = self .batch_size .unwrap_or_else(crate::config::get_default_batch_size); + let document_length = index_document_length(&self.model_id); // Suppress stderr during model loading to hide CoreML's harmless // "Context leak detected" warnings on macOS. @@ -1334,13 +1433,16 @@ impl IndexBuilder { // panic hook and prints it to the restored stderr before resuming, // so panics inside the suppressed region remain visible. let model = crate::stderr::with_suppressed_stderr(|| { - Colbert::builder(&self.model_path) + let mut builder = Colbert::builder(&self.model_path) .with_quantized(self.quantized) .with_parallel(num_sessions) .with_batch_size(batch) .with_dynamic_batch(self.dynamic_batch) - .with_execution_provider(execution_provider) - .build() + .with_execution_provider(execution_provider); + if let Some(length) = model_document_length(document_length) { + builder = builder.with_document_length(length); + } + builder.build() }) .context("Failed to load ColBERT model")?; @@ -1381,15 +1483,19 @@ impl IndexBuilder { .parallel_sessions .unwrap_or_else(|| default_index_parallel_sessions(ExecutionProvider::Cpu, usize::MAX)); let batch = crate::config::DEFAULT_BATCH_SIZE_CPU; + let document_length = index_document_length(&self.model_id); let model = crate::stderr::with_suppressed_stderr(|| { - Colbert::builder(&self.model_path) + let mut builder = Colbert::builder(&self.model_path) .with_quantized(self.quantized) .with_parallel(num_sessions) .with_batch_size(batch) .with_dynamic_batch(false) - .with_execution_provider(ExecutionProvider::Cpu) - .build() + .with_execution_provider(ExecutionProvider::Cpu); + if let Some(length) = model_document_length(document_length) { + builder = builder.with_document_length(length); + } + builder.build() }) .context("Failed to load ColBERT model for CPU fallback")?; @@ -2004,7 +2110,11 @@ impl IndexBuilder { // the encoding pipeline, which rebuilds their FTS5 entries. delete_files_from_index_no_fts_rebuild(index_path, &files_changed)?; - let sorted_units = prepare_units_for_encoding(&new_units, index_chunk_size); + let sorted_units = prepare_units_for_encoding( + &new_units, + index_chunk_size, + index_document_length(&self.model_id), + ); let was_interrupted = self.run_encoding_pipeline( &sorted_units, index_chunk_size, @@ -2287,17 +2397,14 @@ impl IndexBuilder { // Encode in file-coherent batches of ~BUILD_CHECKPOINT_UNITS units, committing state // after each batch so interruptions keep finished work. // - // Binary builds use a single batch instead: the first committed batch - // creates the index and every later batch appends via update_append, - // which binary indexes reject. One batch keeps the whole build on the - // atomic initial-create path (encoded sign bits are ~4x smaller than - // residual codes, so holding them in memory is cheap); the cost is - // that an interrupted binary build restarts from zero. - let checkpoint_units = if self.binary { - usize::MAX - } else { - BUILD_CHECKPOINT_UNITS - }; + // A fresh build uses a single batch so the initial-create path writes the index once. + // Repeated update_append calls rebuild index-wide IVF data, making a new index + // quadratic in its number of checkpoints. A resumed build keeps the smaller + // checkpoints so an interruption only loses its current batch. + // + // Binary builds also require one batch because they cannot be appended to. + let checkpoint_units = + build_checkpoint_units(self.binary, index_dir.join("metadata.json").exists()); let encode_pb = ProgressBar::new(total_units as u64); encode_pb.set_style( ProgressStyle::default_bar() @@ -2410,7 +2517,11 @@ impl IndexBuilder { delete_files_from_index_no_fts_rebuild(index_path, batch_files)?; self.ensure_model_created(batch_units.len())?; let pool_factor = self.resolve_pool_factor(batch_units.len()); - let sorted_units = prepare_units_for_encoding(batch_units, index_chunk_size); + let sorted_units = prepare_units_for_encoding( + batch_units, + index_chunk_size, + index_document_length(&self.model_id), + ); let was_interrupted = self.run_encoding_pipeline( &sorted_units, index_chunk_size, @@ -2774,7 +2885,11 @@ impl IndexBuilder { // new versions; we do a single rebuild after encoding completes. delete_files_from_index_no_fts_rebuild(index_path, &plan.changed)?; - let sorted_units = prepare_units_for_encoding(&new_units, index_chunk_size); + let sorted_units = prepare_units_for_encoding( + &new_units, + index_chunk_size, + index_document_length(&self.model_id), + ); let pipeline_interrupted = self.run_encoding_pipeline( &sorted_units, index_chunk_size, @@ -3213,7 +3328,11 @@ impl IndexBuilder { // Compute effective pool factor based on batch size let pool_factor = self.resolve_pool_factor(units.len()); - let sorted_units = prepare_units_for_encoding(units, index_chunk_size); + let sorted_units = prepare_units_for_encoding( + units, + index_chunk_size, + index_document_length(&self.model_id), + ); self.ensure_model_created(units.len())?; let was_interrupted = self.run_encoding_pipeline( &sorted_units, @@ -4597,6 +4716,75 @@ fn prompt_large_index_confirmation(num_units: usize) -> bool { mod tests { use super::*; + #[test] + fn test_fresh_build_uses_single_write() { + assert_eq!(build_checkpoint_units(false, false), usize::MAX); + } + + #[test] + fn test_resumed_build_keeps_checkpoints() { + assert_eq!(build_checkpoint_units(false, true), BUILD_CHECKPOINT_UNITS); + } + + #[test] + fn test_binary_build_uses_single_write() { + assert_eq!(build_checkpoint_units(true, false), usize::MAX); + assert_eq!(build_checkpoint_units(true, true), usize::MAX); + } + + #[test] + fn test_default_model_uses_adaptive_index_document_length() { + assert_eq!( + resolve_index_document_length(crate::model::DEFAULT_MODEL, None), + IndexDocumentLength::Adaptive + ); + assert_eq!( + resolve_index_document_length("other/model", None), + IndexDocumentLength::Native + ); + assert_eq!( + resolve_index_document_length(crate::model::DEFAULT_MODEL, Some("256")), + IndexDocumentLength::Fixed(256) + ); + assert_eq!( + resolve_index_document_length(crate::model::DEFAULT_MODEL, Some("0")), + IndexDocumentLength::Native + ); + assert_eq!( + resolve_index_document_length(crate::model::DEFAULT_MODEL, Some("invalid")), + IndexDocumentLength::Adaptive + ); + } + + #[test] + fn test_adaptive_document_length_uses_structure() { + let mut unit = CodeUnit::new( + "example".to_string(), + PathBuf::from("src/example.rs"), + 1, + 3, + Language::Rust, + UnitType::Function, + None, + ); + assert_eq!( + code_unit_document_length(&unit, IndexDocumentLength::Adaptive), + 64 + ); + + unit.has_branches = true; + assert_eq!( + code_unit_document_length(&unit, IndexDocumentLength::Adaptive), + 128 + ); + + unit.complexity = 10; + assert_eq!( + code_unit_document_length(&unit, IndexDocumentLength::Adaptive), + 512 + ); + } + /// The mtime fast path in `compute_update_plan` must skip content hashing /// for files whose stored mtime is unchanged. The stored hash is wrong on /// purpose: if the plan still reports the file as unchanged, hashing was diff --git a/colgrep/src/index/state.rs b/colgrep/src/index/state.rs index efeffdc4..3c6d615e 100644 --- a/colgrep/src/index/state.rs +++ b/colgrep/src/index/state.rs @@ -12,7 +12,7 @@ use super::paths::get_state_path; /// metadata schema). Bump ONLY for incompatible changes: a mismatch discards /// the index and re-embeds the entire project on the next run. Routine CLI /// releases must NOT bump this. -pub const INDEX_FORMAT_VERSION: u32 = 2; +pub const INDEX_FORMAT_VERSION: u32 = 3; #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct IndexState { diff --git a/colgrep/src/parser/mod.rs b/colgrep/src/parser/mod.rs index 15a4b0b2..067aa00e 100644 --- a/colgrep/src/parser/mod.rs +++ b/colgrep/src/parser/mod.rs @@ -157,10 +157,18 @@ pub fn extract_units(path: &Path, source: &str, lang: Language) -> Vec } let mut parser = Parser::new(); - if parser - .set_language(&get_tree_sitter_language(lang)) - .is_err() + // The JavaScript grammar treats Flow type annotations as syntax errors and can + // fragment one function into branch-sized pseudo-functions. The TSX grammar + // accepts the same typed JavaScript shape (with or without JSX), preserving the + // coherent outer function for files that opt into Flow. + let parser_language = if lang == Language::JavaScript + && source.lines().take(20).any(|line| line.contains("@flow")) { + tree_sitter_typescript::LANGUAGE_TSX.into() + } else { + get_tree_sitter_language(lang) + }; + if parser.set_language(&parser_language).is_err() { return Vec::new(); } @@ -184,6 +192,7 @@ pub fn extract_units(path: &Path, source: &str, lang: Language) -> Vec lang, &mut units, None, + false, &file_imports, 0, max_depth, @@ -220,6 +229,7 @@ fn extract_from_node( lang: Language, units: &mut Vec, parent_class: Option<&str>, + inside_function: bool, file_imports: &[String], depth: usize, max_depth: usize, @@ -293,6 +303,7 @@ fn extract_from_node( lang, units, Some(&class_name), + inside_function, file_imports, depth + 1, max_depth, @@ -305,7 +316,14 @@ fn extract_from_node( } } // Check if this is a top-level constant/static declaration (only at module level) - else if parent_class.is_none() && is_constant_node(kind, lang) { + else if parent_class.is_none() + && (!inside_function + || !matches!( + lang, + Language::JavaScript | Language::TypeScript | Language::Vue | Language::Svelte + )) + && is_constant_node(kind, lang) + { if let Some(unit) = extract_constant(node, path, lines, bytes, lang, file_imports) { units.push(unit); } @@ -323,6 +341,7 @@ fn extract_from_node( lang, units, parent_class, + inside_function || is_function_node(kind, lang), file_imports, depth + 1, max_depth, diff --git a/colgrep/src/parser/tests/test_javascript.rs b/colgrep/src/parser/tests/test_javascript.rs index e454ace7..a2b5b64a 100644 --- a/colgrep/src/parser/tests/test_javascript.rs +++ b/colgrep/src/parser/tests/test_javascript.rs @@ -24,6 +24,43 @@ function greet(name) { assert_eq!(text, expected); } +#[test] +fn test_flow_exported_function_stays_coherent() { + let source = r#"/** + * @flow + */ +export default function formatConsoleArguments( + maybeMessage: any, + ...inputArgs: $ReadOnlyArray +): $ReadOnlyArray { + if (inputArgs.length === 0) { + return [maybeMessage]; + } + return inputArgs; +}"#; + let units = parse(source, Language::JavaScript, "test.js"); + let function = get_unit_by_name(&units, "formatConsoleArguments").unwrap(); + + assert_eq!(function.line, 1); + assert_eq!(function.end_line, 12); + assert_eq!(function.complexity, 3); + assert!(!units.iter().any(|unit| unit.name == "if")); +} + +#[test] +fn test_function_locals_are_not_duplicate_constant_units() { + let source = r#"const moduleLevel = 1; +function calculate() { + const local = 2; + return moduleLevel + local; +}"#; + let units = parse(source, Language::JavaScript, "test.js"); + + assert!(get_unit_by_name(&units, "moduleLevel").is_some()); + assert!(get_unit_by_name(&units, "calculate").is_some()); + assert!(get_unit_by_name(&units, "local").is_none()); +} + #[test] fn test_arrow_function() { let source = r#"const add = (a, b) => { diff --git a/next-plaid-onnx/src/lib.rs b/next-plaid-onnx/src/lib.rs index acf1b76b..4373987f 100644 --- a/next-plaid-onnx/src/lib.rs +++ b/next-plaid-onnx/src/lib.rs @@ -761,6 +761,11 @@ pub struct PreparedDocumentBatch { struct TokenizedDocument { ids: Vec, type_ids: Vec, + max_length: usize, +} + +fn encoding_worker_count(batch_count: usize, session_count: usize) -> usize { + batch_count.min(session_count).max(1) } impl PreparedDocumentBatch { @@ -1144,14 +1149,40 @@ impl Colbert { pub fn tokenize_documents_in_batches( &self, documents: &[&str], + ) -> Result> { + self.tokenize_documents_in_batches_with_limits(documents, &[]) + } + + /// Tokenize documents using a separate maximum sequence length for each document. + /// + /// An empty `max_lengths` slice uses the model's configured document length for + /// every document. Otherwise it must contain one entry per document. + pub fn tokenize_documents_in_batches_with_limits( + &self, + documents: &[&str], + max_lengths: &[usize], ) -> Result> { if documents.is_empty() { return Ok(Vec::new()); } + if !max_lengths.is_empty() && max_lengths.len() != documents.len() { + anyhow::bail!( + "Expected one document length per document, got {} lengths for {} documents", + max_lengths.len(), + documents.len() + ); + } let processed_texts = preprocess_texts(&self.config, documents); - let tokenized = tokenize_processed_texts_individually(&self.tokenizer, &processed_texts)?; - let truncate_limit = self.config.document_length.saturating_sub(1); + let mut tokenized = + tokenize_processed_texts_individually(&self.tokenizer, &processed_texts)?; + for (index, document) in tokenized.iter_mut().enumerate() { + document.max_length = max_lengths + .get(index) + .copied() + .unwrap_or(self.config.document_length) + .clamp(2, self.config.document_length); + } let use_gpu_batch_modes = !matches!(self.requested_execution_provider, ExecutionProvider::Cpu); let use_dynamic_batch = self.dynamic_batch && use_gpu_batch_modes; @@ -1196,7 +1227,12 @@ impl Colbert { // input order in the returned embeddings. let prepared_lengths: Vec = tokenized .iter() - .map(|doc| doc.ids.len().min(truncate_limit) + 1) + .map(|doc| { + doc.ids + .len() + .min(doc.max_length.saturating_sub(1)) + .saturating_add(1) + }) .collect(); let mut items: Vec<(usize, usize, TokenizedDocument)> = prepared_lengths .into_iter() @@ -1308,41 +1344,60 @@ impl Colbert { all_embeddings } else { let cancel_ref = cancel; - let results: Vec>>> = std::thread::scope(|scope| { - let mut handles = Vec::with_capacity(prepared_batches.len()); - - for (i, prepared_batch) in prepared_batches.into_iter().enumerate() { - let session_idx = i % self.sessions.len(); - let session_mutex = &self.sessions[session_idx]; - let config = &self.config; - let skiplist_ids = &self.skiplist_ids; - - handles.push(scope.spawn(move || { - // Skip not-yet-started batches once cancelled; already-running - // forward passes finish (one `session.run`), so the whole call - // returns within ~one batch of a Ctrl+C. - if cancel_ref.map(|c| c()).unwrap_or(false) { - anyhow::bail!("encoding cancelled"); - } - let mut session = session_mutex.lock().unwrap(); - encode_prepared_batch_with_session( - &mut session, - config, - skiplist_ids, - prepared_batch, - ) - })); - } + let worker_count = encoding_worker_count(prepared_batches.len(), self.sessions.len()); + let work: Mutex> = + Mutex::new(prepared_batches.into_iter().enumerate().collect()); + let mut results: Vec<(usize, Vec>)> = + std::thread::scope(|scope| -> Result<_> { + let mut handles = Vec::with_capacity(worker_count); + + for session_idx in 0..worker_count { + let session_mutex = &self.sessions[session_idx]; + let config = &self.config; + let skiplist_ids = &self.skiplist_ids; + let work = &work; + + handles.push(scope.spawn(move || -> Result<_> { + let mut session = session_mutex.lock().unwrap(); + let mut encoded_batches = Vec::new(); + loop { + let Some((batch_index, prepared_batch)) = + work.lock().unwrap().pop_front() + else { + break; + }; + // Skip not-yet-started batches once cancelled; already-running + // forward passes finish (one `session.run`), so the whole call + // returns within ~one batch of a Ctrl+C. + if cancel_ref.map(|c| c()).unwrap_or(false) { + anyhow::bail!("encoding cancelled"); + } + let embeddings = encode_prepared_batch_with_session( + &mut session, + config, + skiplist_ids, + prepared_batch, + )?; + encoded_batches.push((batch_index, embeddings)); + } + Ok(encoded_batches) + })); + } - handles - .into_iter() - .map(|handle| handle.join().unwrap()) - .collect() - }); + let mut results = Vec::new(); + for handle in handles { + let worker_results = handle + .join() + .map_err(|_| anyhow::anyhow!("encoding worker thread panicked"))??; + results.extend(worker_results); + } + Ok(results) + })?; + results.sort_unstable_by_key(|(batch_index, _)| *batch_index); let mut all_embeddings = Vec::new(); - for result in results { - all_embeddings.extend(result?); + for (_, batch_embeddings) in results { + all_embeddings.extend(batch_embeddings); } all_embeddings }; @@ -1732,6 +1787,7 @@ fn tokenize_processed_texts_individually( Ok(TokenizedDocument { ids: encoding.get_ids()[..real_len].to_vec(), type_ids: encoding.get_type_ids()[..real_len].to_vec(), + max_length: usize::MAX, }) }) .collect::>() @@ -1909,9 +1965,10 @@ fn prepare_batch_from_tokenized_documents( })?, }; - let truncate_limit = max_length.saturating_sub(1); let mut batch_max_len = 0usize; for doc in &batch_docs { + let max_length = doc.max_length.min(max_length); + let truncate_limit = max_length.saturating_sub(1); let effective_len = if doc.ids.len() > truncate_limit { max_length } else { @@ -1946,6 +2003,8 @@ fn prepare_batch_from_tokenized_documents( for (row_idx, doc) in batch_docs.into_iter().enumerate() { let row_start = row_idx * batch_max_len; let real_len = doc.ids.len().max(1); + let max_length = doc.max_length.min(max_length); + let truncate_limit = max_length.saturating_sub(1); let (content_prefix_len, keep_sep) = if real_len > truncate_limit { (truncate_limit.saturating_sub(1), true) } else { @@ -2320,6 +2379,13 @@ fn pool_embeddings_hierarchical( mod tests { use super::*; + #[test] + fn test_encoding_worker_count_is_bounded() { + assert_eq!(encoding_worker_count(1024, 14), 14); + assert_eq!(encoding_worker_count(2, 14), 2); + assert_eq!(encoding_worker_count(0, 14), 1); + } + // ========================================================================= // ColbertConfig tests // =========================================================================