1mod backends;
2
3use crate::chunking::{ChunkingAlgorithm, SlidingWindowChunker};
4use chrono::Duration;
5use indexmap::IndexMap;
6use nohash_hasher::IntSet;
7use rand::prelude::*;
8use rayon::prelude::*;
9use std::borrow::Cow;
10use std::collections::{HashMap, HashSet};
11use std::path::Path;
12use std::sync::atomic::{AtomicUsize, Ordering};
13use std::sync::mpsc;
14use std::sync::{Arc, Mutex};
15use std::thread;
16
17use crate::config::{
18 ChunkingStrategy, NegativeStrategy, SamplerConfig, Selector, TextRecipe, TripletRecipe,
19};
20use crate::constants::sampler::AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME;
21use crate::constants::sampler::{
22 ANCHOR_POSITIVE_SWAP_MASK, EPOCH_SEED_OFFSET, EXHAUSTION_RETRY_LIMIT, NEG_REASON_SAME_RECORD,
23 NEG_REASON_WRONG_ARTICLE, NEG_REASON_WRONG_DATE, NEG_REASON_WRONG_QA, PREFETCHER_SOURCE_ID,
24 PREFETCHER_STOPPED_REASON, RECIPE_LABEL_TEXT, RECIPE_LABEL_TRIPLETS,
25 RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER, ROLE_LABEL_ANCHOR, ROLE_LABEL_CONTEXT,
26 SAME_SELECTOR_PAIR_RETRY_LIMIT,
27};
28use crate::data::{
29 ChunkView, DataRecord, PairLabel, RecordChunk, RecordSection, SampleBatch, SamplePair,
30 SampleTriplet, SectionRole, TextBatch, TextSample, TripletBatch,
31};
32use crate::epoch::EpochTracker;
33use crate::errors::SamplerError;
34use crate::hash::{derive_epoch_seed, stable_hash_key, stable_hash_key_section, stable_hash_str};
35use crate::ingestion::IngestionManager;
36use crate::metadata::{META_FIELD_DATE, MetadataKey};
37use crate::metrics::{chunk_proximity_score, window_index_proximity};
38use crate::source::DataSource;
39use crate::splits::{
40 EpochStateStore, PersistedSamplerState, SamplerStateStore, SplitLabel, SplitStore,
41};
42use crate::tokenizer::{Tokenizer, WhitespaceTokenizer};
43use crate::types::{RecipeKey, RecordKey, SourceId};
44use crate::utils::platform_newline;
45
46#[derive(Debug, Clone)]
59struct DeterministicRng {
61 state: u64,
62}
63
64impl DeterministicRng {
65 fn new(seed: u64) -> Self {
66 Self { state: seed }
67 }
68
69 fn from_state(state: u64) -> Self {
70 Self { state }
71 }
72
73 fn state(&self) -> u64 {
74 self.state
75 }
76
77 fn next_u64_internal(&mut self) -> u64 {
78 let mut z = self.state.wrapping_add(0x9E3779B97F4A7C15);
79 self.state = z;
80 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
81 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
82 z ^ (z >> 31)
83 }
84}
85
86impl rand::RngCore for DeterministicRng {
87 fn next_u32(&mut self) -> u32 {
88 self.next_u64_internal() as u32
89 }
90
91 fn next_u64(&mut self) -> u64 {
92 self.next_u64_internal()
93 }
94
95 fn fill_bytes(&mut self, dest: &mut [u8]) {
96 let mut offset = 0;
97 while offset < dest.len() {
98 let value = self.next_u64_internal();
99 let bytes = value.to_le_bytes();
100 let remaining = dest.len() - offset;
101 let copy_len = remaining.min(bytes.len());
102 dest[offset..offset + copy_len].copy_from_slice(&bytes[..copy_len]);
103 offset += copy_len;
104 }
105 }
106}
107
108pub fn chunk_weight(strategy: &ChunkingStrategy, chunk: &RecordChunk) -> f32 {
117 let floor = strategy.chunk_weight_floor;
118 let trust = chunk.quality.trust.clamp(0.0, 1.0);
119 let base = match &chunk.view {
120 ChunkView::Window { index, .. } => window_index_proximity(*index),
121 ChunkView::SummaryFallback { weight, .. } => *weight,
122 };
123 (base * trust).max(floor)
124}
125
126pub trait Sampler {
128 #[deprecated(
133 note = "samples all sources uniformly; call next_pair_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
134 )]
135 fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
136 self.next_pair_batch_with_weights(split, &HashMap::new())
137 }
138 #[deprecated(
143 note = "samples all sources uniformly; call next_text_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
144 )]
145 fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
146 self.next_text_batch_with_weights(split, &HashMap::new())
147 }
148 #[deprecated(
153 note = "samples all sources uniformly; call next_triplet_batch_with_weights(split, &weights) with an explicit per-source weight map to honor a data mixture"
154 )]
155 fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
156 self.next_triplet_batch_with_weights(split, &HashMap::new())
157 }
158 fn next_pair_batch_with_weights(
160 &self,
161 split: SplitLabel,
162 weights: &HashMap<SourceId, f32>,
163 ) -> Result<SampleBatch, SamplerError>;
164 fn next_text_batch_with_weights(
166 &self,
167 split: SplitLabel,
168 weights: &HashMap<SourceId, f32>,
169 ) -> Result<TextBatch, SamplerError>;
170 fn next_triplet_batch_with_weights(
172 &self,
173 split: SplitLabel,
174 weights: &HashMap<SourceId, f32>,
175 ) -> Result<TripletBatch, SamplerError>;
176 fn save_sampler_state(&self, _save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
180 Ok(())
181 }
182}
183
184pub struct BatchPrefetcher<T> {
186 receiver: Option<mpsc::Receiver<Result<T, SamplerError>>>,
187 handle: Option<thread::JoinHandle<()>>,
188 stats: Arc<PrefetcherStats>,
189}
190
191#[derive(Default)]
192struct PrefetcherStats {
194 queued: AtomicUsize,
195 produced: AtomicUsize,
196 errors: AtomicUsize,
197}
198
199impl<T: Send + 'static> BatchPrefetcher<T> {
200 fn new<F>(capacity: usize, mut producer: F) -> Self
201 where
202 F: FnMut() -> Result<T, SamplerError> + Send + 'static,
203 {
204 let (sender, receiver) = mpsc::sync_channel(capacity.max(1));
205 let stats = Arc::new(PrefetcherStats::default());
206 let stats_thread = Arc::clone(&stats);
207 let handle = thread::spawn(move || {
208 loop {
209 let result = producer();
210 if result.is_err() {
211 stats_thread.errors.fetch_add(1, Ordering::Relaxed);
212 }
213 if sender.send(result).is_err() {
214 return;
215 }
216 stats_thread.queued.fetch_add(1, Ordering::Relaxed);
217 stats_thread.produced.fetch_add(1, Ordering::Relaxed);
218 }
219 });
220 Self {
221 receiver: Some(receiver),
222 handle: Some(handle),
223 stats,
224 }
225 }
226
227 pub fn next(&self) -> Result<T, SamplerError> {
229 let receiver = self
230 .receiver
231 .as_ref()
232 .ok_or_else(|| SamplerError::SourceUnavailable {
233 source_id: PREFETCHER_SOURCE_ID.into(),
234 reason: PREFETCHER_STOPPED_REASON.into(),
235 })?;
236 let result = receiver.recv().unwrap_or_else(|_| {
237 Err(SamplerError::SourceUnavailable {
238 source_id: PREFETCHER_SOURCE_ID.into(),
239 reason: PREFETCHER_STOPPED_REASON.into(),
240 })
241 });
242 self.stats
243 .queued
244 .try_update(Ordering::Relaxed, Ordering::Relaxed, |value| {
245 Some(value.saturating_sub(1))
246 })
247 .ok();
248 result
249 }
250
251 pub fn queue_len(&self) -> usize {
253 self.stats.queued.load(Ordering::Relaxed)
254 }
255
256 pub fn produced_count(&self) -> usize {
258 self.stats.produced.load(Ordering::Relaxed)
259 }
260
261 pub fn error_count(&self) -> usize {
263 self.stats.errors.load(Ordering::Relaxed)
264 }
265}
266
267impl<T> Drop for BatchPrefetcher<T> {
268 fn drop(&mut self) {
269 self.receiver.take();
270 if let Some(handle) = self.handle.take() {
271 let _ = handle.join();
272 }
273 }
274}
275
276pub struct TripletSampler<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
279 inner: Mutex<TripletSamplerInner<S>>,
280}
281
282struct TripletSamplerInner<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
284 config: SamplerConfig,
286 chunker: Arc<dyn ChunkingAlgorithm>,
288 split_store: Arc<S>,
290 ingestion: IngestionManager,
292 records: IndexMap<RecordKey, Arc<DataRecord>>,
296 source_deficits: HashMap<SourceId, f32>,
305 rng: DeterministicRng,
307 triplet_recipes: Vec<TripletRecipe>,
309 text_recipes: Vec<TextRecipe>,
311 source_triplet_recipes: HashMap<SourceId, Vec<TripletRecipe>>,
313 long_section_counts: HashMap<SourceId, usize>,
319 source_text_recipes: HashMap<SourceId, Vec<TextRecipe>>,
321 using_config_triplet_recipes: bool,
323 using_config_text_recipes: bool,
325 last_observed_ingest: u64,
327 epoch_tracker: EpochTracker,
329 chunk_cursors: HashMap<(RecordKey, usize), usize>,
331 role_cursors: HashMap<(RecordKey, String), usize>,
333 negative_backend: Box<dyn backends::NegativeBackend>,
335 chunk_index: HashMap<RecordKey, RecordKey>,
337 source_order: Vec<SourceId>,
339 source_cycle_idx: usize,
341 source_state_loaded: bool,
343 ingestion_cursors_loaded: bool,
345 source_state_dirty: bool,
347 source_record_indices: HashMap<SourceId, Vec<RecordKey>>,
349 source_record_cursors: HashMap<SourceId, usize>,
351 triplet_recipe_rr_idx: usize,
353 emitted_text_hashes: HashMap<RecordKey, HashSet<u64>>,
362 text_recipe_rr_idx: usize,
364 epoch: u64,
366
367 source_wrapped: HashMap<SourceId, bool>,
369
370 last_sync_version: u64,
372
373 pool_generation: u64,
378
379 cached_records_by_split: Option<HashMap<SplitLabel, Vec<RecordKey>>>,
383 cached_records_by_split_gen: u64,
385
386 last_force_refresh_generation: u64,
392
393 split_labels: HashMap<RecordKey, SplitLabel>,
396}
397
398impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSamplerInner<S> {
399 fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
400 Self::new_with_chunker(config, split_store, Arc::new(SlidingWindowChunker))
401 }
402
403 fn new_with_chunker(
404 config: SamplerConfig,
405 split_store: Arc<S>,
406 chunker: Arc<dyn ChunkingAlgorithm>,
407 ) -> Self {
408 let buffer_size = config.ingestion_max_records.max(config.batch_size).max(2);
409 let using_config_triplet_recipes = !config.recipes.is_empty();
410 let using_config_text_recipes = !config.text_recipes.is_empty();
411 let triplet_recipes = if using_config_triplet_recipes {
412 config.recipes.clone()
413 } else {
414 Vec::new()
415 };
416 let text_recipes = if using_config_text_recipes {
417 config.text_recipes.clone()
418 } else if !triplet_recipes.is_empty() {
419 Self::build_derived_text_recipes(&triplet_recipes)
420 } else {
421 Vec::new()
422 };
423 let ingestion = IngestionManager::new(buffer_size, config.clone());
424 let epoch_backend = Some(Arc::clone(&split_store) as Arc<dyn EpochStateStore>);
425 let epoch_tracker = EpochTracker::new(
426 true,
427 epoch_backend,
428 derive_epoch_seed(config.seed, EPOCH_SEED_OFFSET),
429 );
430 let mut sampler = Self {
431 rng: DeterministicRng::new(config.seed),
432 config,
433 chunker,
434 split_store,
435 ingestion,
436 records: IndexMap::new(),
437 triplet_recipes,
438 text_recipes,
439 source_triplet_recipes: HashMap::new(),
440 long_section_counts: HashMap::new(),
441 source_text_recipes: HashMap::new(),
442 using_config_triplet_recipes,
443 using_config_text_recipes,
444 last_observed_ingest: 0,
445 epoch_tracker,
446 chunk_cursors: HashMap::new(),
447 role_cursors: HashMap::new(),
448 negative_backend: {
449 #[cfg(feature = "bm25-mining")]
450 {
451 Box::new(backends::Bm25Backend::new())
452 }
453 #[cfg(not(feature = "bm25-mining"))]
454 {
455 Box::new(backends::DefaultBackend)
456 }
457 },
458 chunk_index: HashMap::new(),
459 source_order: Vec::new(),
460 source_cycle_idx: 0,
461 source_deficits: HashMap::new(),
462 source_state_loaded: false,
463 ingestion_cursors_loaded: false,
464 source_state_dirty: false,
465 source_record_indices: HashMap::new(),
466 source_record_cursors: HashMap::new(),
467 emitted_text_hashes: HashMap::new(),
468 triplet_recipe_rr_idx: 0,
469 text_recipe_rr_idx: 0,
470 epoch: 0,
471 source_wrapped: HashMap::new(),
472 last_sync_version: 0,
473 last_force_refresh_generation: 0,
474 pool_generation: 0,
475 cached_records_by_split: None,
476 cached_records_by_split_gen: 0,
477 split_labels: HashMap::new(),
478 };
479 if !sampler.using_config_text_recipes {
480 sampler.rebuild_derived_text_recipes();
481 }
482 sampler
483 }
484
485 fn text_recipes(&self) -> &[TextRecipe] {
486 &self.text_recipes
487 }
488
489 fn epoch_seed(&self) -> u64 {
492 derive_epoch_seed(self.config.seed, self.epoch)
493 }
494
495 fn get_or_insert_split_label(&mut self, key: &RecordKey) -> Result<SplitLabel, SamplerError> {
497 if let Some(&label) = self.split_labels.get(key) {
498 return Ok(label);
499 }
500 let label = match self.split_store.label_for(key) {
501 Some(label) => label,
502 None => self.split_store.ensure(key.clone())?,
503 };
504 self.split_labels.insert(key.clone(), label);
505 Ok(label)
506 }
507
508 fn register_source(
509 &mut self,
510 source: Box<dyn DataSource + 'static>,
511 ) -> Result<(), SamplerError> {
512 let source_id = source.id().to_string();
513 if !self.using_config_triplet_recipes {
514 let triplets = source.default_triplet_recipes();
515 if !triplets.is_empty() {
516 self.source_triplet_recipes
517 .insert(source_id.clone(), triplets.clone());
518 if !self.using_config_text_recipes {
519 let derived = Self::build_derived_text_recipes(&triplets);
520 self.source_text_recipes
521 .insert(source_id.clone(), derived.clone());
522 self.extend_text_recipes_unique(&derived);
523 }
524 }
525 }
526 self.ingestion.register_source(source)?;
527 Ok(())
528 }
529
530 fn set_epoch(&mut self, epoch: u64) -> Result<(), SamplerError> {
531 self.epoch_tracker.ensure_loaded()?;
532 self.epoch_tracker.force_epoch(epoch);
533 self.epoch = epoch;
534 self.ingestion.set_epoch(epoch);
535 self.ingestion.reset_stream_cursors();
536 self.ingestion.reset_epoch_step();
539 self.source_record_cursors.clear();
540 self.source_cycle_idx = 0;
541 for source in &self.source_order {
542 self.source_wrapped.insert(source.clone(), false);
543 }
544 self.rebuild_source_index()?;
545 self.source_state_dirty = self.source_order.len() > 1;
546 Ok(())
547 }
548
549 fn next_chunk_from_pool(
550 &mut self,
551 key: &RecordKey,
552 section_idx: usize,
553 pool: Vec<RecordChunk>,
554 ) -> Option<RecordChunk> {
555 if pool.is_empty() {
556 return None;
557 }
558 let map_key = (key.clone(), section_idx);
559 if !self.chunk_cursors.contains_key(&map_key) {
560 let cursor_key = stable_hash_key_section(self.epoch_seed(), key, section_idx);
566 let start = (cursor_key as usize) % pool.len();
567 self.chunk_cursors.insert(map_key.clone(), start);
568 }
569 let cursor = self.chunk_cursors.entry(map_key).or_insert(0);
570 if *cursor >= pool.len() {
571 *cursor = 0;
572 }
573 let chunk = pool.get(*cursor).cloned();
574 *cursor = (*cursor + 1) % pool.len();
575 chunk
576 }
577
578 fn prune_cursor_state(&mut self) {
579 if self.chunk_cursors.is_empty()
580 && self.role_cursors.is_empty()
581 && self.negative_backend.cursors_empty()
582 {
583 return;
584 }
585 let valid_keys: HashSet<RecordKey> = self.records.keys().cloned().collect();
586 self.chunk_cursors
587 .retain(|(record_key, _), _| valid_keys.contains(record_key));
588 self.role_cursors
589 .retain(|(record_key, _), _| valid_keys.contains(record_key));
590 self.negative_backend.prune_cursors(&valid_keys);
591 }
592
593 fn remove_id_from_source_index(&mut self, source: &SourceId, record_key: &RecordKey) {
599 if let Some(ids) = self.source_record_indices.get_mut(source)
600 && let Some(pos) = ids.iter().position(|key| key == record_key)
601 {
602 ids.remove(pos);
603 }
604 }
605
606 fn insert_id_into_source_index(
612 &mut self,
613 source: SourceId,
614 record_key: RecordKey,
615 label: SplitLabel,
616 allowed: &HashSet<SplitLabel>,
617 ) {
618 if !allowed.contains(&label) {
619 return;
620 }
621 let seed = self.epoch_seed();
622 let key = stable_hash_key(seed, &record_key);
623 let pos = match self.source_record_indices.get(&source) {
627 Some(ids) => ids.partition_point(|k| stable_hash_key(seed, k) <= key),
628 None => 0,
629 };
630 self.source_record_indices
631 .entry(source)
632 .or_default()
633 .insert(pos, record_key);
634 }
635
636 fn increment_long_section_count(&mut self, source: &SourceId) {
638 *self.long_section_counts.entry(source.clone()).or_default() += 1;
639 }
640
641 fn decrement_long_section_count(&mut self, source: &SourceId) {
644 let remove = if let Some(count) = self.long_section_counts.get_mut(source) {
645 *count = count.saturating_sub(1);
646 *count == 0
647 } else {
648 false
649 };
650 if remove {
651 self.long_section_counts.remove(source);
652 }
653 }
654
655 fn rebuild_chunk_index(&mut self) {
656 self.chunk_index.clear();
657 for record in self.records.values() {
658 let key = RecordKey::from(record);
659 self.chunk_index.insert(key.clone(), key);
660 }
661 }
662
663 fn rebuild_source_index(&mut self) -> Result<(), SamplerError> {
664 self.source_record_indices.clear();
665 let allowed = self.allowed_target_splits();
666 let allowed_set: HashSet<SplitLabel> = allowed.into_iter().collect();
667 let record_keys: Vec<RecordKey> = self.records.values().map(RecordKey::of_arc).collect();
669 for record_key in record_keys {
670 let label = self.get_or_insert_split_label(&record_key)?;
671 if !allowed_set.contains(&label) {
672 continue;
673 }
674 self.source_record_indices
675 .entry(record_key.source.clone())
676 .or_default()
677 .push(record_key);
678 }
679
680 let shuffle_seed = self.epoch_seed();
681 for record_keys in self.source_record_indices.values_mut() {
682 record_keys.sort_by_key(|key| stable_hash_key(shuffle_seed, key));
683 }
684
685 self.source_order = self.source_record_indices.keys().cloned().collect();
686 self.source_order.sort();
687 self.refresh_source_wrapped();
688 self.source_deficits
690 .retain(|s, _| self.source_record_indices.contains_key(s));
691
692 self.source_record_cursors
693 .retain(|source, _| self.source_record_indices.contains_key(source));
694 if self.source_state_loaded {
695 if self.source_order.is_empty() {
696 self.source_cycle_idx = 0;
697 }
698 self.source_state_dirty = self.source_order.len() > 1;
699 }
700 Ok(())
701 }
702
703 fn refresh_source_wrapped(&mut self) {
704 self.source_wrapped.clear();
705 for source in &self.source_order {
706 let len = self
707 .source_record_indices
708 .get(source)
709 .map(|items| items.len())
710 .unwrap_or(0);
711 if len == 0 {
712 self.source_wrapped.insert(source.clone(), false);
713 continue;
714 }
715 let cursor = self.source_record_cursors.get(source).copied().unwrap_or(0);
716 let wrapped = cursor > 0 && cursor % len == 0;
717 self.source_wrapped.insert(source.clone(), wrapped);
718 }
719 }
720
721 fn shuffled_source_cycle(&self, cycle: u64) -> Vec<SourceId> {
722 let mut sources = self.source_order.clone();
723 let seed = self.epoch_seed() ^ cycle;
724 sources.sort_by_key(|source| stable_hash_str(seed, source));
725 sources
726 }
727
728 fn weights_are_selective(weights: &HashMap<SourceId, f32>) -> bool {
731 if weights.is_empty() {
732 return false;
733 }
734 let first = *weights.values().next().unwrap();
735 weights.values().any(|&w| w != first)
736 }
737
738 fn weight_quantum(weights: &HashMap<SourceId, f32>, source: &str) -> f32 {
742 weights.get(source).copied().unwrap_or(1.0).max(0.0)
743 }
744
745 fn select_weighted_source(
751 &mut self,
752 candidates: &[SourceId],
753 weights: &HashMap<SourceId, f32>,
754 ) -> Option<SourceId> {
755 let eligible: Vec<&SourceId> = candidates
756 .iter()
757 .filter(|s| {
758 self.source_record_indices
759 .get(*s)
760 .map(|v| !v.is_empty())
761 .unwrap_or(false)
762 })
763 .collect();
764 if eligible.is_empty() {
765 return None;
766 }
767 let mut total = 0.0f32;
768 for source in &eligible {
769 let q = Self::weight_quantum(weights, source);
770 total += q;
771 *self.source_deficits.entry((*source).clone()).or_insert(0.0) += q;
772 }
773 if total <= 0.0 {
774 return None;
775 }
776 let mut best: Option<(f32, SourceId)> = None;
777 for source in &eligible {
778 let d = self.source_deficits.get(*source).copied().unwrap_or(0.0);
779 if best.as_ref().map(|(bd, _)| d > *bd).unwrap_or(true) {
780 best = Some((d, (*source).clone()));
781 }
782 }
783 let (_, source) = best?;
784 *self
785 .source_deficits
786 .get_mut(&source)
787 .expect("deficit credited above") -= total;
788 Some(source)
789 }
790
791 fn ensure_source_state(&mut self) -> Result<(), SamplerError> {
792 if self.source_state_loaded {
793 return Ok(());
794 }
795 let persisted = self.split_store.load_sampler_state()?;
796 self.source_cycle_idx = persisted
797 .as_ref()
798 .map(|state| state.source_cycle_idx as usize)
799 .unwrap_or(0);
800 if let Some(state) = persisted {
801 for (source, cursor) in state.source_record_cursors {
802 if self.source_record_indices.contains_key(&source) {
803 self.source_record_cursors.insert(source, cursor as usize);
804 }
805 }
806 self.epoch = state.epoch;
807 self.ingestion.set_epoch(state.epoch);
808 self.rng = DeterministicRng::from_state(state.rng_state);
809 self.triplet_recipe_rr_idx = state.triplet_recipe_rr_idx as usize;
810 self.text_recipe_rr_idx = state.text_recipe_rr_idx as usize;
811 }
812 self.refresh_source_wrapped();
813 self.source_state_loaded = true;
814 self.source_state_dirty = true;
815 Ok(())
816 }
817
818 fn persist_source_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
819 if !self.source_state_loaded {
820 return Ok(());
821 }
822 let state = PersistedSamplerState {
823 source_cycle_idx: self.source_cycle_idx as u64,
824 source_record_cursors: self
825 .source_record_cursors
826 .iter()
827 .map(|(source, cursor)| (source.clone(), *cursor as u64))
828 .collect(),
829 epoch: self.epoch,
830 epoch_step: self.ingestion.epoch_step(),
831 rng_state: self.rng.state(),
832 triplet_recipe_rr_idx: self.triplet_recipe_rr_idx as u64,
833 text_recipe_rr_idx: self.text_recipe_rr_idx as u64,
834 source_stream_cursors: self.ingestion.snapshot_cursors(),
835 };
836 self.split_store.save_sampler_state(&state, save_to)?;
837 self.source_state_dirty = false;
838 Ok(())
839 }
840
841 fn rebuild_derived_text_recipes(&mut self) {
842 if self.using_config_text_recipes {
843 return;
844 }
845 if self.triplet_recipes.is_empty() {
846 self.text_recipes.clear();
847 } else {
848 self.text_recipes = Self::build_derived_text_recipes(&self.triplet_recipes);
849 }
850 }
851
852 fn extend_text_recipes_unique(&mut self, recipes: &[TextRecipe]) {
853 for recipe in recipes {
854 if self
855 .text_recipes
856 .iter()
857 .any(|existing| existing.name == recipe.name)
858 {
859 continue;
860 }
861 self.text_recipes.push(recipe.clone());
862 }
863 }
864
865 fn configured_triplet_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TripletRecipe] {
866 if self.using_config_triplet_recipes {
867 return &self.triplet_recipes;
868 }
869 self.source_triplet_recipes
870 .get(source)
871 .map(|recipes| recipes.as_slice())
872 .unwrap_or(&[])
873 }
874
875 fn contains_auto_chunk_pair_recipe(recipes: &[TripletRecipe]) -> bool {
877 recipes
878 .iter()
879 .any(|recipe| recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME)
880 }
881
882 fn source_supports_chunk_pair_recipe(&self, source: &str) -> bool {
883 if self.config.chunking.max_window_tokens == 0 {
884 return false;
885 }
886 self.long_section_counts
887 .get(source)
888 .is_some_and(|&count| count > 0)
889 }
890
891 fn should_auto_inject_chunk_pair_recipe(
897 &self,
898 source: &str,
899 recipes: &[TripletRecipe],
900 ) -> bool {
901 self.source_supports_chunk_pair_recipe(source)
902 && !Self::contains_auto_chunk_pair_recipe(recipes)
903 }
904
905 fn source_chunk_pair_recipe() -> TripletRecipe {
916 TripletRecipe {
917 name: Cow::Borrowed(AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME),
918 anchor: Selector::Role(SectionRole::Context),
919 positive_selector: Selector::Role(SectionRole::Context),
920 negative_selector: Selector::Role(SectionRole::Context),
921 negative_strategy: NegativeStrategy::WrongArticle,
922 weight: 1.0,
923 instruction: None,
924 allow_same_anchor_positive: false,
925 }
926 }
927
928 fn resolve_source_triplet_plan(&self, source: &str) -> (Vec<TripletRecipe>, bool) {
940 let mut recipes = self.configured_triplet_recipes_for_source(source).to_vec();
941 let mut auto_injected = false;
942 if self.should_auto_inject_chunk_pair_recipe(source, &recipes) {
943 recipes.push(Self::source_chunk_pair_recipe());
944 auto_injected = true;
945 }
946 (recipes, auto_injected)
947 }
948
949 #[cfg(test)]
950 fn triplet_recipes_for_source(&self, source: &str) -> Vec<TripletRecipe> {
951 self.resolve_source_triplet_plan(source).0
952 }
953
954 fn triplet_recipe_count_for_source(&self, source: &str) -> usize {
955 let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
956 recipes.len()
957 }
958
959 fn text_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TextRecipe] {
960 if self.using_config_text_recipes || self.using_config_triplet_recipes {
961 return &self.text_recipes;
962 }
963 self.source_text_recipes
964 .get(source)
965 .map(|recipes| recipes.as_slice())
966 .unwrap_or(&[])
967 }
968
969 fn recipe_order_weighted_shuffled(
980 &mut self,
981 weights: &[f32],
982 rng: &mut DeterministicRng,
983 ) -> Vec<usize> {
984 weighted_recipe_order(weights, rng)
985 }
986
987 fn recipe_order_weighted_cycled(
993 &mut self,
994 weights: &[f32],
995 rr_idx: usize,
996 rng: &mut DeterministicRng,
997 ) -> Vec<usize> {
998 let base = self.recipe_order_weighted_shuffled(weights, rng);
999 if base.is_empty() {
1000 return base;
1001 }
1002 let start = rr_idx % base.len();
1003 let mut order = Vec::with_capacity(base.len());
1004 order.extend_from_slice(&base[start..]);
1005 order.extend_from_slice(&base[..start]);
1006 order
1007 }
1008
1009 fn text_recipe_order_weighted_shuffled(
1013 &mut self,
1014 weights: &[f32],
1015 rng: &mut DeterministicRng,
1016 ) -> Vec<usize> {
1017 weighted_recipe_order(weights, rng)
1018 }
1019
1020 fn text_recipe_order_weighted_cycled(
1021 &mut self,
1022 weights: &[f32],
1023 rr_idx: usize,
1024 rng: &mut DeterministicRng,
1025 ) -> Vec<usize> {
1026 let base = self.text_recipe_order_weighted_shuffled(weights, rng);
1027 if base.is_empty() {
1028 return base;
1029 }
1030 let start = rr_idx % base.len();
1031 let mut order = Vec::with_capacity(base.len());
1032 order.extend_from_slice(&base[start..]);
1033 order.extend_from_slice(&base[..start]);
1034 order
1035 }
1036
1037 fn allowed_target_splits(&self) -> Vec<SplitLabel> {
1038 self.config.allowed_splits.clone()
1039 }
1040
1041 fn ensure_split_allowed(&self, split: SplitLabel) -> Result<(), SamplerError> {
1042 let allowed = self.allowed_target_splits();
1043 if allowed.contains(&split) {
1044 return Ok(());
1045 }
1046 Err(SamplerError::Configuration(format!(
1047 "requested split {:?} is not in allowed_splits {:?}",
1048 split, allowed
1049 )))
1050 }
1051
1052 fn ensure_split_has_records(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
1053 self.refresh_records_by_split_cache()?;
1054 let has_records = self
1055 .cached_records_by_split
1056 .as_ref()
1057 .and_then(|map| map.get(&target_split))
1058 .is_some_and(|records| !records.is_empty());
1059 if has_records {
1060 return Ok(());
1061 }
1062 Err(SamplerError::Exhausted(
1063 "no records available for target split".into(),
1064 ))
1065 }
1066
1067 fn refresh_records_by_split_cache(&mut self) -> Result<(), SamplerError> {
1071 if self.cached_records_by_split_gen == self.pool_generation
1072 && self.cached_records_by_split.is_some()
1073 {
1074 return Ok(());
1075 }
1076 let mut map: HashMap<SplitLabel, Vec<RecordKey>> = HashMap::new();
1077 let entries: Vec<(RecordKey, RecordKey)> = self
1079 .chunk_index
1080 .iter()
1081 .map(|(k, v)| (k.clone(), v.clone()))
1082 .collect();
1083 for (chunk_key, record_key) in entries {
1084 if !self.records.contains_key(&record_key) {
1085 continue;
1086 }
1087 let label = self.get_or_insert_split_label(&record_key)?;
1088 map.entry(label).or_default().push(chunk_key);
1089 }
1090 self.cached_records_by_split = Some(map);
1091 self.cached_records_by_split_gen = self.pool_generation;
1092 Ok(())
1093 }
1094
1095 #[cfg(test)]
1099 fn records_by_split(&mut self) -> Result<HashMap<SplitLabel, Vec<RecordKey>>, SamplerError> {
1100 self.refresh_records_by_split_cache()?;
1101 Ok(self.cached_records_by_split.clone().unwrap_or_default())
1102 }
1103
1104 fn choose_anchor_record(
1105 &mut self,
1106 source: Option<&str>,
1107 split: SplitLabel,
1108 ) -> Option<Arc<DataRecord>> {
1109 if let Some(source) = source {
1110 let record_ids = self.source_record_indices.get(source)?;
1111 if record_ids.is_empty() {
1112 return None;
1113 }
1114 let len = record_ids.len();
1115 let mut cursor = *self.source_record_cursors.get(source).unwrap_or(&0);
1116 let cycle = cursor / len;
1117 let offset_seed = self.epoch_seed() ^ (cycle as u64);
1118 let offset = (stable_hash_str(offset_seed, source) as usize) % len;
1119 let mut wrapped = false;
1120 let mut selected_key: Option<RecordKey> = None;
1121
1122 for _ in 0..len {
1123 let pos = (cursor % len).saturating_add(offset) % len;
1124 let record_key = &record_ids[pos];
1125 cursor = cursor.saturating_add(1);
1126 if cursor.is_multiple_of(len) {
1127 wrapped = true;
1128 }
1129 if self.split_labels.get(record_key) == Some(&split) {
1130 selected_key = Some(record_key.clone());
1131 break;
1132 }
1133 }
1134
1135 self.source_record_cursors
1136 .insert(source.to_string(), cursor);
1137 if wrapped {
1138 self.mark_source_wrapped(source);
1139 }
1140
1141 selected_key.and_then(|key| self.records.get(&key).map(Arc::clone))
1142 } else {
1143 while let Some(chunk_key) = self.epoch_tracker.next_record(split) {
1144 if let Some(record_key) = self.chunk_index.get(&chunk_key)
1145 && let Some(record) = self.records.get(record_key)
1146 {
1147 return Some(Arc::clone(record));
1148 }
1149 }
1150 None
1151 }
1152 }
1153
1154 fn save_sampler_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
1155 if self.epoch_tracker.is_enabled() {
1156 self.epoch_tracker.persist()?;
1157 }
1158 self.persist_source_state(save_to)?;
1159 Ok(())
1160 }
1161
1162 fn mark_source_wrapped(&mut self, source: &str) {
1163 self.source_wrapped.insert(source.to_string(), true);
1164 if self.source_order.is_empty() {
1165 return;
1166 }
1167 let all_wrapped = self
1168 .source_order
1169 .iter()
1170 .all(|name| self.source_wrapped.get(name).copied().unwrap_or(false));
1171 if all_wrapped {
1172 self.advance_epoch();
1173 }
1174 }
1175
1176 fn advance_epoch(&mut self) {
1177 self.epoch = self.epoch.saturating_add(1);
1184 self.ingestion.set_epoch(self.epoch);
1185 self.source_record_cursors.clear();
1192 self.source_cycle_idx = 0;
1193 for source in &self.source_order {
1196 self.source_wrapped.insert(source.clone(), false);
1197 }
1198 let _ = self.rebuild_source_index();
1199 self.source_state_dirty = self.source_order.len() > 1;
1200 }
1201
1202 fn select_temporal_neighbor(
1203 &'_ self,
1204 record: &DataRecord,
1205 offset_days: i32,
1206 ) -> Option<Arc<DataRecord>> {
1207 let target = record.created_at + Duration::days(offset_days.into());
1208 let key = record.taxonomy.first().cloned();
1209 let record_split = self.split_store.label_for(&RecordKey::from(record))?;
1210 self.records
1211 .values()
1212 .filter(|candidate| {
1213 RecordKey::from(*candidate) != RecordKey::from(record)
1214 && self
1215 .split_store
1216 .label_for(&RecordKey::from(*candidate))
1217 .map(|label| label == record_split)
1218 .unwrap_or(false)
1219 && (candidate.source == record.source
1220 || key
1221 .as_ref()
1222 .zip(candidate.taxonomy.first())
1223 .map(|(a, b)| a == b)
1224 .unwrap_or(false))
1225 })
1226 .min_by_key(|candidate| (candidate.created_at - target).num_seconds().abs())
1227 .cloned()
1228 }
1229
1230 fn select_negative_record(
1231 &self,
1232 anchor_record: &DataRecord,
1233 strategy: &NegativeStrategy,
1234 anchor_query_text: Option<&str>,
1235 rng: &mut dyn rand::RngCore,
1236 ) -> Option<(Arc<DataRecord>, bool)> {
1237 let anchor_split = self
1238 .split_store
1239 .label_for(&RecordKey::from(anchor_record))?;
1240
1241 let in_anchor_split = |candidate: &DataRecord| {
1242 self.split_store
1243 .label_for(&RecordKey::from(candidate))
1244 .map(|label| label == anchor_split)
1245 .unwrap_or(false)
1246 };
1247
1248 match strategy {
1249 NegativeStrategy::WrongArticle => {
1250 let anchor_date =
1251 taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1252 let mut same_date: Vec<Arc<DataRecord>> = self
1253 .records
1254 .values()
1255 .filter(|candidate| {
1256 candidate.source == anchor_record.source
1257 && candidate.id != anchor_record.id
1258 && in_anchor_split(candidate)
1259 })
1260 .filter(|candidate| {
1261 anchor_date
1262 .as_deref()
1263 .zip(taxonomy_value(candidate, META_FIELD_DATE))
1264 .map(|(a, b)| a == b)
1265 .unwrap_or(false)
1266 })
1267 .cloned()
1268 .collect();
1269 if same_date.is_empty() {
1270 same_date = self
1271 .records
1272 .values()
1273 .filter(|candidate| {
1274 candidate.source == anchor_record.source
1275 && candidate.id != anchor_record.id
1276 && in_anchor_split(candidate)
1277 })
1278 .cloned()
1279 .collect();
1280 }
1281 if !same_date.is_empty() {
1282 return self.negative_backend.choose_negative(
1283 anchor_record,
1284 anchor_split,
1285 same_date,
1286 false,
1287 anchor_query_text,
1288 rng,
1289 );
1290 }
1291 let pool = self
1292 .records
1293 .values()
1294 .filter(|candidate| {
1295 candidate.id != anchor_record.id && in_anchor_split(candidate)
1296 })
1297 .cloned()
1298 .collect::<Vec<_>>();
1299 self.negative_backend.choose_negative(
1300 anchor_record,
1301 anchor_split,
1302 pool,
1303 true,
1304 anchor_query_text,
1305 rng,
1306 )
1307 }
1308 NegativeStrategy::WrongPublicationDate => {
1309 let anchor_date =
1310 taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1311 let pool: Vec<Arc<DataRecord>> = self
1312 .records
1313 .values()
1314 .filter(|candidate| {
1315 candidate.source == anchor_record.source
1316 && candidate.id != anchor_record.id
1317 && in_anchor_split(candidate)
1318 })
1319 .filter(|candidate| {
1320 match (
1321 anchor_date.as_deref(),
1322 taxonomy_value(candidate, META_FIELD_DATE),
1323 ) {
1324 (Some(anchor), Some(candidate_date)) => anchor != candidate_date,
1325 (Some(_), None) => true,
1326 (None, Some(_)) => true,
1327 (None, None) => false,
1328 }
1329 })
1330 .cloned()
1331 .collect();
1332 if pool.is_empty() {
1333 let fallback_pool = self
1336 .records
1337 .values()
1338 .filter(|candidate| {
1339 candidate.id != anchor_record.id && in_anchor_split(candidate)
1340 })
1341 .cloned()
1342 .collect::<Vec<_>>();
1343
1344 return self.negative_backend.choose_negative(
1345 anchor_record,
1346 anchor_split,
1347 fallback_pool,
1348 true,
1349 anchor_query_text,
1350 rng,
1351 );
1352 }
1353
1354 self.negative_backend.choose_negative(
1355 anchor_record,
1356 anchor_split,
1357 pool,
1358 false,
1359 anchor_query_text,
1360 rng,
1361 )
1362 }
1363 NegativeStrategy::QuestionAnswerMismatch => {
1364 let pool: Vec<Arc<DataRecord>> = self
1365 .records
1366 .values()
1367 .filter(|candidate| {
1368 candidate.source == anchor_record.source
1369 && candidate.id != anchor_record.id
1370 && in_anchor_split(candidate)
1371 })
1372 .cloned()
1373 .collect();
1374 if pool.is_empty() {
1375 let fallback_pool = self
1378 .records
1379 .values()
1380 .filter(|candidate| {
1381 candidate.id != anchor_record.id && in_anchor_split(candidate)
1382 })
1383 .cloned()
1384 .collect::<Vec<_>>();
1385
1386 return self.negative_backend.choose_negative(
1387 anchor_record,
1388 anchor_split,
1389 fallback_pool,
1390 true,
1391 anchor_query_text,
1392 rng,
1393 );
1394 }
1395
1396 self.negative_backend.choose_negative(
1397 anchor_record,
1398 anchor_split,
1399 pool,
1400 false,
1401 anchor_query_text,
1402 rng,
1403 )
1404 }
1405 NegativeStrategy::SameRecord => {
1406 Some((Arc::new(anchor_record.clone()), false))
1412 }
1413 }
1414 }
1415
1416 fn is_auto_chunk_pair_recipe(recipe: &TripletRecipe) -> bool {
1418 recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME
1419 }
1420
1421 fn select_anchor_positive_pair(
1425 &mut self,
1426 record: &DataRecord,
1427 anchor_selector: &Selector,
1428 positive_selector: &Selector,
1429 enforce_window_pair: bool,
1430 ) -> Option<(RecordChunk, RecordChunk)> {
1431 let mut anchor_chunk = self.select_chunk(record, anchor_selector)?;
1432 let mut positive_chunk = self.select_chunk(record, positive_selector)?;
1433 if anchor_selector == positive_selector {
1434 let mut retries = 0usize;
1435 while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1436 && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1437 {
1438 let Some(redraw_anchor) = self.select_chunk(record, anchor_selector) else {
1439 break;
1440 };
1441 let Some(redraw_positive) = self.select_chunk(record, positive_selector) else {
1442 break;
1443 };
1444 anchor_chunk = redraw_anchor;
1445 positive_chunk = redraw_positive;
1446 retries += 1;
1447 }
1448 if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1449 return None;
1450 }
1451 }
1452 Some((anchor_chunk, positive_chunk))
1453 }
1454
1455 fn select_distinct_window_pair_for_auto_recipe(
1465 &mut self,
1466 recipe: &TripletRecipe,
1467 record: &DataRecord,
1468 ) -> Option<(RecordChunk, RecordChunk)> {
1469 if recipe.anchor != recipe.positive_selector {
1470 return None;
1471 }
1472 self.select_anchor_positive_pair(record, &recipe.anchor, &recipe.positive_selector, true)
1473 }
1474
1475 fn record_has_at_least_two_window_chunks_for_selector(
1480 &self,
1481 record: &DataRecord,
1482 selector: &Selector,
1483 ) -> bool {
1484 let section_indices: Vec<usize> = match selector {
1485 Selector::Role(role) => record
1486 .sections
1487 .iter()
1488 .enumerate()
1489 .filter(|(_, section)| roles_match(role, §ion.role))
1490 .map(|(idx, _)| idx)
1491 .collect(),
1492 Selector::Paragraph(idx) => {
1493 if *idx < record.sections.len() {
1494 vec![*idx]
1495 } else {
1496 Vec::new()
1497 }
1498 }
1499 Selector::Random => (0..record.sections.len()).collect(),
1500 Selector::TemporalOffset(_) => return false,
1501 };
1502
1503 let mut window_count = 0usize;
1504 for section_idx in section_indices {
1505 let Some(section) = record.sections.get(section_idx) else {
1506 continue;
1507 };
1508 let chunks = self.materialize_chunks(record, section_idx, section);
1509 window_count += chunks
1510 .iter()
1511 .filter(|chunk| matches!(chunk.view, ChunkView::Window { .. }))
1512 .count();
1513 if window_count >= 2 {
1514 return true;
1515 }
1516 }
1517 false
1518 }
1519
1520 fn make_triplet_with_unique_negative(
1532 &mut self,
1533 recipe: &TripletRecipe,
1534 record: &DataRecord,
1535 rng: &mut DeterministicRng,
1536 seen: &mut HashSet<(RecordKey, RecordKey, RecordKey)>,
1537 seen_negatives: &mut IntSet<u64>,
1538 ) -> Option<SampleTriplet> {
1539 let (mut anchor_chunk, mut positive_chunk) = if Self::is_auto_chunk_pair_recipe(recipe) {
1541 if !self.record_has_at_least_two_window_chunks_for_selector(record, &recipe.anchor) {
1542 return None;
1543 }
1544 self.select_distinct_window_pair_for_auto_recipe(recipe, record)?
1545 } else {
1546 self.select_anchor_positive_pair(
1547 record,
1548 &recipe.anchor,
1549 &recipe.positive_selector,
1550 false,
1551 )?
1552 };
1553 let anchor_raw_text = anchor_chunk.text.clone();
1554 self.decorate_chunk(record, &mut anchor_chunk, rng);
1555 self.decorate_chunk(record, &mut positive_chunk, rng);
1556
1557 let (final_anchor, final_positive) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0 {
1559 (positive_chunk, anchor_chunk)
1560 } else {
1561 (anchor_chunk, positive_chunk)
1562 };
1563 if !recipe.allow_same_anchor_positive && final_anchor.text == final_positive.text {
1564 return None;
1565 }
1566
1567 let mut accepted: Option<(RecordChunk, f32, bool)> = None;
1569 for _ in 0..NEG_DEDUP_RETRIES {
1570 let (negative_record, fallback_used) = self.select_negative_record(
1571 record,
1572 &recipe.negative_strategy,
1573 Some(&anchor_raw_text),
1574 rng,
1575 )?;
1576 let mut negative_chunk =
1577 match self.select_chunk(&negative_record, &recipe.negative_selector) {
1578 Some(c) => c,
1579 None => continue,
1580 };
1581 self.decorate_chunk(&negative_record, &mut negative_chunk, rng);
1582 if negative_chunk.text == final_positive.text
1583 || negative_chunk.text == final_anchor.text
1584 {
1585 continue;
1586 }
1587 let key = (
1588 RecordKey::of_chunk(&final_anchor),
1589 RecordKey::of_chunk(&final_positive),
1590 RecordKey::of_chunk(&negative_chunk),
1591 );
1592 let neg_h = negative_text_hash(&negative_chunk);
1593 if seen.contains(&key) || seen_negatives.contains(&neg_h) {
1595 continue;
1596 }
1597 let chunk_weight =
1598 self.triplet_chunk_weight(recipe, &final_anchor, &final_positive, &negative_chunk);
1599 let weight = recipe.weight * chunk_weight;
1600 seen.insert(key);
1601 seen_negatives.insert(neg_h);
1602 accepted = Some((negative_chunk, weight, fallback_used));
1603 break;
1604 }
1605 let (negative_chunk, weight, fallback_used) = accepted?;
1606 let recipe_name = if fallback_used {
1607 format!("{}_fallback_same_split", recipe.name)
1608 } else {
1609 recipe.name.to_string()
1610 };
1611 Some(SampleTriplet {
1612 recipe: recipe_name,
1613 anchor: final_anchor,
1614 positive: final_positive,
1615 negative: negative_chunk,
1616 weight,
1617 instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1618 })
1619 }
1620
1621 fn make_text_sample_for_split(
1622 &mut self,
1623 recipe: &TextRecipe,
1624 source: Option<&str>,
1625 split: SplitLabel,
1626 rng: &mut DeterministicRng,
1627 ) -> Option<TextSample> {
1628 let record = self.choose_anchor_record(source, split)?;
1629 let mut chunk = self.select_chunk(&record, &recipe.selector)?;
1630 self.decorate_chunk(&record, &mut chunk, rng);
1631 let weight = recipe.weight * self.chunk_weight(&chunk);
1632 Some(TextSample {
1633 recipe: recipe.name.to_string(),
1634 chunk,
1635 weight,
1636 instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1637 })
1638 }
1639
1640 fn chunk_weight(&self, chunk: &RecordChunk) -> f32 {
1641 chunk_weight(&self.config.chunking, chunk)
1642 }
1643
1644 fn triplet_chunk_weight(
1645 &self,
1646 recipe: &TripletRecipe,
1647 anchor: &RecordChunk,
1648 positive: &RecordChunk,
1649 negative: &RecordChunk,
1650 ) -> f32 {
1651 let floor = self.config.chunking.chunk_weight_floor;
1652 let negative_weight = negative.quality.trust.clamp(0.0, 1.0).max(floor);
1653 if Self::is_auto_chunk_pair_recipe(recipe) {
1654 let pair_trust = ((anchor.quality.trust.clamp(0.0, 1.0)
1657 + positive.quality.trust.clamp(0.0, 1.0))
1658 / 2.0)
1659 .clamp(0.0, 1.0);
1660 let pair_weight = (chunk_proximity_score(anchor, positive) * pair_trust).max(floor);
1661 return (pair_weight + pair_weight + negative_weight) / 3.0;
1663 }
1664 let pair_proximity = chunk_proximity_score(anchor, positive);
1667 let anchor_weight = (self.chunk_weight(anchor) * pair_proximity).max(floor);
1668 let positive_weight = (self.chunk_weight(positive) * pair_proximity).max(floor);
1669 (anchor_weight + positive_weight + negative_weight) / 3.0
1670 }
1671
1672 fn decorate_chunk(
1673 &mut self,
1674 record: &DataRecord,
1675 chunk: &mut RecordChunk,
1676 rng: &mut DeterministicRng,
1677 ) {
1678 chunk.kvp_meta = record
1679 .meta_prefix
1680 .as_ref()
1681 .map(|s| s.all_metadata())
1682 .unwrap_or_default();
1683 if let Some(spec) = record.meta_prefix.as_ref()
1684 && let Some(prefix) = spec.sample(rng)
1685 {
1686 let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1687 let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1688 let total_tokens = prefix_tokens.len() + body_tokens.len();
1689 let max_window = self.config.chunking.max_window_tokens;
1690 if max_window > 0 && total_tokens > max_window {
1691 if prefix_tokens.len() >= max_window {
1692 chunk.text = prefix_tokens
1693 .into_iter()
1694 .take(max_window)
1695 .collect::<Vec<_>>()
1696 .join(" ");
1697 chunk.tokens_estimate = max_window;
1698 } else {
1699 let remaining = max_window - prefix_tokens.len();
1700 let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1701 chunk.text =
1702 format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1703 chunk.tokens_estimate = max_window;
1704 }
1705 } else {
1706 chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1707 chunk.tokens_estimate = total_tokens;
1708 }
1709 }
1710 }
1711
1712 fn select_chunk_parallel(
1716 &self,
1717 record: &DataRecord,
1718 selector: &Selector,
1719 rng: &mut DeterministicRng,
1720 ) -> Option<RecordChunk> {
1721 match selector {
1722 Selector::Role(role) => self.select_role_parallel(record, role, rng),
1723 Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1724 let pool = self.materialize_chunks(record, *idx, section);
1725 if pool.is_empty() {
1726 return None;
1727 }
1728 let i = rng.random_range(0..pool.len());
1729 pool.into_iter().nth(i)
1730 }),
1731 Selector::TemporalOffset(offset) => self
1732 .select_temporal_neighbor(record, *offset)
1733 .and_then(|neighbor| {
1734 self.select_role_parallel(&neighbor, &SectionRole::Context, rng)
1735 }),
1736 Selector::Random => {
1737 if record.sections.is_empty() {
1738 return None;
1739 }
1740 let idx = rng.random_range(0..record.sections.len());
1741 record.sections.get(idx).and_then(|section| {
1742 let pool = self.materialize_chunks(record, idx, section);
1743 if pool.is_empty() {
1744 return None;
1745 }
1746 let i = rng.random_range(0..pool.len());
1747 pool.into_iter().nth(i)
1748 })
1749 }
1750 }
1751 }
1752
1753 fn select_role_parallel(
1755 &self,
1756 record: &DataRecord,
1757 role: &SectionRole,
1758 rng: &mut DeterministicRng,
1759 ) -> Option<RecordChunk> {
1760 let indices: Vec<usize> = record
1761 .sections
1762 .iter()
1763 .enumerate()
1764 .filter(|(_, s)| roles_match(role, &s.role))
1765 .map(|(i, _)| i)
1766 .collect();
1767 if indices.is_empty() {
1768 return None;
1769 }
1770 let start = rng.random_range(0..indices.len());
1771 for offset in 0..indices.len() {
1772 let section_idx = indices[(start + offset) % indices.len()];
1773 let section = &record.sections[section_idx];
1774 let pool = self.materialize_chunks(record, section_idx, section);
1775 if !pool.is_empty() {
1776 let i = rng.random_range(0..pool.len());
1777 return pool.into_iter().nth(i);
1778 }
1779 }
1780 None
1781 }
1782
1783 fn decorate_chunk_parallel(
1785 &self,
1786 record: &DataRecord,
1787 chunk: &mut RecordChunk,
1788 rng: &mut DeterministicRng,
1789 ) {
1790 chunk.kvp_meta = record
1791 .meta_prefix
1792 .as_ref()
1793 .map(|s| s.all_metadata())
1794 .unwrap_or_default();
1795 if let Some(spec) = record.meta_prefix.as_ref()
1796 && let Some(prefix) = spec.sample(rng)
1797 {
1798 let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1799 let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1800 let total_tokens = prefix_tokens.len() + body_tokens.len();
1801 let max_window = self.config.chunking.max_window_tokens;
1802 if max_window > 0 && total_tokens > max_window {
1803 if prefix_tokens.len() >= max_window {
1804 chunk.text = prefix_tokens
1805 .into_iter()
1806 .take(max_window)
1807 .collect::<Vec<_>>()
1808 .join(" ");
1809 chunk.tokens_estimate = max_window;
1810 } else {
1811 let remaining = max_window - prefix_tokens.len();
1812 let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1813 chunk.text =
1814 format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1815 chunk.tokens_estimate = max_window;
1816 }
1817 } else {
1818 chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1819 chunk.tokens_estimate = total_tokens;
1820 }
1821 }
1822 }
1823
1824 fn select_anchor_positive_parallel(
1826 &self,
1827 record: &DataRecord,
1828 anchor_selector: &Selector,
1829 positive_selector: &Selector,
1830 enforce_window_pair: bool,
1831 rng: &mut DeterministicRng,
1832 ) -> Option<(RecordChunk, RecordChunk)> {
1833 let anchor_chunk = self.select_chunk_parallel(record, anchor_selector, rng)?;
1834 let mut positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1835 if anchor_selector == positive_selector {
1836 let mut retries = 0usize;
1837 while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1838 && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1839 {
1840 positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1841 retries += 1;
1842 }
1843 if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1844 return None;
1845 }
1846 }
1847 Some((anchor_chunk, positive_chunk))
1849 }
1850
1851 fn select_anchor_positive_for_recipe(
1854 &self,
1855 recipe: &TripletRecipe,
1856 anchor_record: &DataRecord,
1857 rng: &mut DeterministicRng,
1858 ) -> Option<(RecordChunk, RecordChunk, String)> {
1859 if Self::is_auto_chunk_pair_recipe(recipe) {
1860 if !self
1861 .record_has_at_least_two_window_chunks_for_selector(anchor_record, &recipe.anchor)
1862 {
1863 return None;
1864 }
1865 let mut anchor_chunk =
1866 self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1867 let mut positive_chunk =
1868 self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1869 let mut tries = 0usize;
1870 while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, true) {
1871 tries += 1;
1872 if tries >= SAME_SELECTOR_PAIR_RETRY_LIMIT {
1873 return None;
1874 }
1875 anchor_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1876 positive_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1877 }
1878 let anchor_raw_text = anchor_chunk.text.clone();
1879 self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1880 self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1881 return Some((anchor_chunk, positive_chunk, anchor_raw_text));
1882 }
1883 let (mut anchor_chunk, mut positive_chunk) = self.select_anchor_positive_parallel(
1884 anchor_record,
1885 &recipe.anchor,
1886 &recipe.positive_selector,
1887 false,
1888 rng,
1889 )?;
1890 let anchor_raw_text = anchor_chunk.text.clone();
1891 self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1892 self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1893 Some((anchor_chunk, positive_chunk, anchor_raw_text))
1894 }
1895
1896 fn select_chunk(&mut self, record: &DataRecord, selector: &Selector) -> Option<RecordChunk> {
1897 match selector {
1898 Selector::Role(role) => self.select_by_role(record, role),
1899 Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1900 let pool = self.materialize_chunks(record, *idx, section);
1901 self.next_chunk_from_pool(&RecordKey::from(record), *idx, pool)
1902 }),
1903 Selector::TemporalOffset(offset) => self
1904 .select_temporal_neighbor(record, *offset)
1905 .and_then(|neighbor| self.select_by_role(&neighbor, &SectionRole::Context)),
1906 Selector::Random => {
1907 if record.sections.is_empty() {
1908 return None;
1909 }
1910 let idx = self.rng.random_range(0..record.sections.len());
1911 record.sections.get(idx).and_then(|section| {
1912 let pool = self.materialize_chunks(record, idx, section);
1913 self.next_chunk_from_pool(&RecordKey::from(record), idx, pool)
1914 })
1915 }
1916 }
1917 }
1918
1919 fn select_by_role(&mut self, record: &DataRecord, role: &SectionRole) -> Option<RecordChunk> {
1920 let indices: Vec<usize> = record
1921 .sections
1922 .iter()
1923 .enumerate()
1924 .filter(|(_, section)| roles_match(role, §ion.role))
1925 .map(|(idx, _)| idx)
1926 .collect();
1927 if indices.is_empty() {
1928 return None;
1929 }
1930 let key = role_cursor_key(&RecordKey::from(record), role);
1931 let start_offset = self
1932 .role_cursors
1933 .get(&key)
1934 .and_then(|last_idx| indices.iter().position(|idx| idx == last_idx))
1935 .map(|pos| (pos + 1) % indices.len())
1936 .unwrap_or_else(|| {
1937 let record_key = RecordKey::from(record);
1943 let seed_key = format!("{}::{}", record_key.id, role_label(role));
1944 (stable_hash_str(self.epoch_seed(), &seed_key) as usize) % indices.len()
1945 });
1946 for offset in 0..indices.len() {
1947 let section_idx = indices[(start_offset + offset) % indices.len()];
1948 let section = &record.sections[section_idx];
1949 let pool = self.materialize_chunks(record, section_idx, section);
1950 if let Some(chunk) =
1951 self.next_chunk_from_pool(&RecordKey::from(record), section_idx, pool)
1952 {
1953 self.role_cursors.insert(key.clone(), section_idx);
1954 return Some(chunk);
1955 }
1956 }
1957 None
1958 }
1959
1960 fn materialize_chunks(
1973 &self,
1974 record: &DataRecord,
1975 section_idx: usize,
1976 section: &RecordSection,
1977 ) -> Vec<RecordChunk> {
1978 let mut chunks =
1979 self.chunker
1980 .materialize(&self.config.chunking, record, section_idx, section);
1981 for chunk in &mut chunks {
1985 chunk.embedding = section.embedding.clone();
1986 chunk.source = record.source.clone();
1987 }
1988 chunks
1989 }
1990
1991 fn build_derived_text_recipes(recipes: &[TripletRecipe]) -> Vec<TextRecipe> {
1992 let mut derived = Vec::new();
1993 for recipe in recipes {
1994 let base = recipe.name.as_ref();
1995 derived.push(TextRecipe {
1996 name: Cow::Owned(format!("{base}_anchor")),
1997 selector: recipe.anchor.clone(),
1998 weight: recipe.weight,
1999 instruction: None,
2000 });
2001 derived.push(TextRecipe {
2002 name: Cow::Owned(format!("{base}_positive")),
2003 selector: recipe.positive_selector.clone(),
2004 weight: recipe.weight,
2005 instruction: None,
2006 });
2007 derived.push(TextRecipe {
2008 name: Cow::Owned(format!("{base}_negative")),
2009 selector: recipe.negative_selector.clone(),
2010 weight: recipe.weight,
2011 instruction: None,
2012 });
2013 }
2014 derived
2015 }
2016
2017 fn record_has_long_anchor_or_context_section(&self, record: &DataRecord) -> bool {
2020 let window = self.config.chunking.max_window_tokens;
2021 if window == 0 {
2022 return false;
2023 }
2024 record.sections.iter().any(|s| {
2025 matches!(s.role, SectionRole::Anchor | SectionRole::Context) && s.token_count > window
2026 })
2027 }
2028 fn sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
2029 if self.last_sync_version == 0 {
2030 self.full_sync_records_from_cache()?;
2031 return Ok(());
2032 }
2033
2034 if self.ingestion.force_refresh_generation() != self.last_force_refresh_generation {
2039 self.last_force_refresh_generation = self.ingestion.force_refresh_generation();
2040 self.full_sync_records_from_cache()?;
2041 return Ok(());
2042 }
2043
2044 let (new_version, added, evicted) = self.ingestion.sync_delta(self.last_sync_version);
2047
2048 if added.is_empty() && evicted.is_empty() {
2049 self.negative_backend.on_sync_start();
2050 self.prune_cursor_state();
2051 self.last_sync_version = new_version;
2052 return Ok(());
2053 }
2054
2055 let mut truly_evicted = false;
2063 if !evicted.is_empty() {
2064 let added_keys: HashSet<RecordKey> = added.iter().map(RecordKey::of_arc).collect();
2065 for evicted_key in &evicted {
2066 if added_keys.contains(evicted_key) {
2067 continue;
2068 }
2069 let Some(record) = self.records.get(evicted_key) else {
2070 continue;
2071 };
2072 let source = record.source.clone();
2073 let is_long = self.record_has_long_anchor_or_context_section(record);
2074 self.records.shift_remove(evicted_key);
2075 self.emitted_text_hashes.remove(evicted_key);
2076 self.split_labels.remove(evicted_key);
2077 self.chunk_index.remove(evicted_key);
2078 self.remove_id_from_source_index(&source, evicted_key);
2079 if is_long {
2080 self.decrement_long_section_count(&source);
2081 }
2082 truly_evicted = true;
2083 }
2084 }
2085
2086 self.negative_backend.on_sync_start();
2087
2088 let mut pool_changed = truly_evicted;
2093 if !added.is_empty() {
2094 let allowed: HashSet<SplitLabel> = self.allowed_target_splits().into_iter().collect();
2095 for record in &added {
2096 let key = RecordKey::from(record);
2103 if self.records.contains_key(&key) {
2104 self.records.insert(key, Arc::clone(record));
2105 continue;
2106 }
2107 pool_changed = true;
2108 let label = self.get_or_insert_split_label(&key)?;
2109 if self.record_has_long_anchor_or_context_section(record) {
2110 self.increment_long_section_count(&record.source);
2111 }
2112 self.records.insert(key.clone(), Arc::clone(record));
2113 self.chunk_index.insert(key.clone(), key.clone());
2114 self.insert_id_into_source_index(record.source.clone(), key, label, &allowed);
2115 }
2116 }
2117
2118 let empty_sources: Vec<SourceId> = self
2123 .source_record_indices
2124 .iter()
2125 .filter(|(_, ids)| ids.is_empty())
2126 .map(|(source, _)| source.clone())
2127 .collect();
2128 let mut order_changed = !empty_sources.is_empty();
2129 for source in &empty_sources {
2130 self.source_record_indices.remove(source);
2131 }
2132 if !order_changed {
2134 for source in self.source_record_indices.keys() {
2135 if !self.source_order.contains(source) {
2136 order_changed = true;
2137 break;
2138 }
2139 }
2140 }
2141 if order_changed {
2142 self.source_order = self.source_record_indices.keys().cloned().collect();
2143 self.source_order.sort();
2144 }
2145 self.source_deficits
2147 .retain(|s, _| self.source_record_indices.contains_key(s));
2148 self.refresh_source_wrapped();
2149 self.source_record_cursors
2150 .retain(|source, _| self.source_record_indices.contains_key(source));
2151 if self.source_state_loaded {
2152 if self.source_order.is_empty() {
2153 self.source_cycle_idx = 0;
2154 }
2155 self.source_state_dirty = self.source_order.len() > 1;
2156 }
2157 if truly_evicted {
2159 self.prune_cursor_state();
2160 }
2161
2162 if pool_changed {
2163 self.pool_generation = self.pool_generation.saturating_add(1);
2164 }
2165 self.last_sync_version = new_version;
2166 Ok(())
2167 }
2168
2169 fn full_sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
2173 let mut snapshot = self.ingestion.all_records_snapshot();
2174 snapshot.sort_by_key(RecordKey::of_arc);
2175
2176 let old_keys: HashSet<RecordKey> = self.records.keys().cloned().collect();
2178 let incoming_keys: HashSet<RecordKey> = snapshot.iter().map(RecordKey::of_arc).collect();
2179 let evicted: HashSet<&RecordKey> = old_keys.difference(&incoming_keys).collect();
2180
2181 self.records.clear();
2182 self.long_section_counts.clear();
2183 for evicted_key in &evicted {
2184 self.emitted_text_hashes.remove(*evicted_key);
2185 self.split_labels.remove(*evicted_key);
2186 }
2187 self.negative_backend.on_sync_start();
2188 for record in snapshot {
2189 let key = RecordKey::from(&record);
2190 let _ = self.get_or_insert_split_label(&key)?;
2191 if self.record_has_long_anchor_or_context_section(&record) {
2192 self.increment_long_section_count(&record.source);
2193 }
2194 self.records.insert(key, Arc::clone(&record));
2195 }
2196 self.prune_cursor_state();
2197 self.rebuild_chunk_index();
2198 self.rebuild_source_index()?;
2199 self.last_sync_version = self.ingestion.global_version();
2200 self.pool_generation = self.pool_generation.saturating_add(1);
2202 Ok(())
2203 }
2204
2205 fn ingest_internal_for_split(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
2206 if !self.ingestion.has_sources() {
2207 return Ok(());
2208 }
2209 if !self.ingestion_cursors_loaded {
2210 if let Some(state) = self.split_store.load_sampler_state()? {
2211 self.ingestion.load_cursors(&state.source_stream_cursors);
2212 self.ingestion.set_epoch(state.epoch);
2213 self.ingestion.set_epoch_step(state.epoch_step);
2214 }
2215 self.ingestion_cursors_loaded = true;
2216 }
2217 if self.ingestion.all_caches_empty() {
2218 self.ingestion.refresh_all();
2219 } else {
2220 self.ingestion.advance(self.config.batch_size);
2221 }
2222 let observed = self.ingestion.total_ingest_count();
2223 if observed == self.last_observed_ingest && !self.records.is_empty() {
2224 return Ok(());
2225 }
2226 self.last_observed_ingest = observed;
2227 self.sync_records_from_cache()?;
2228 let max_window_tokens = self.config.chunking.max_window_tokens;
2229 self.negative_backend.on_records_refreshed(
2230 &self.records,
2231 max_window_tokens,
2232 &|id| self.split_labels.get(id).copied(),
2233 self.ingestion.last_refreshed_sources(),
2234 );
2235 self.epoch_tracker.ensure_loaded()?;
2238 self.refresh_records_by_split_cache()?;
2239 let cached_records_by_split = self
2240 .cached_records_by_split
2241 .as_ref()
2242 .expect("records_by_split cache populated by refresh");
2243 self.epoch_tracker
2244 .reconcile(target_split, cached_records_by_split);
2245 self.ensure_source_state()?;
2246 Ok(())
2247 }
2248
2249 fn ensure_ingestion_cursors_loaded(&mut self) -> Result<(), SamplerError> {
2258 if !self.ingestion_cursors_loaded {
2259 if let Some(state) = self.split_store.load_sampler_state()? {
2260 self.ingestion.load_cursors(&state.source_stream_cursors);
2261 self.ingestion.set_epoch(state.epoch);
2262 self.ingestion.set_epoch_step(state.epoch_step);
2263 }
2264 self.ingestion_cursors_loaded = true;
2265 }
2266 Ok(())
2267 }
2268
2269 #[cfg(test)]
2270 fn ingest_internal(&mut self, split: SplitLabel) -> Result<(), SamplerError> {
2271 self.ingest_internal_for_split(split)
2272 }
2273
2274 fn ingest_internal_with_weights_for_split(
2275 &mut self,
2276 target_split: SplitLabel,
2277 weights: &HashMap<SourceId, f32>,
2278 ) -> Result<(), SamplerError> {
2279 if !self.ingestion.has_sources() {
2280 return Ok(());
2281 }
2282 if !self.ingestion_cursors_loaded {
2283 if let Some(state) = self.split_store.load_sampler_state()? {
2284 self.ingestion.load_cursors(&state.source_stream_cursors);
2285 self.ingestion.set_epoch(state.epoch);
2286 self.ingestion.set_epoch_step(state.epoch_step);
2287 }
2288 self.ingestion_cursors_loaded = true;
2289 }
2290 if self.ingestion.all_caches_empty() {
2291 self.ingestion.refresh_all_with_weights(weights)?;
2292 } else {
2293 self.ingestion
2294 .advance_with_weights(self.config.batch_size, weights)?;
2295 }
2296 let observed = self.ingestion.total_ingest_count();
2297 if observed == self.last_observed_ingest && !self.records.is_empty() {
2298 return Ok(());
2299 }
2300 self.last_observed_ingest = observed;
2301 self.sync_records_from_cache()?;
2302 let max_window_tokens = self.config.chunking.max_window_tokens;
2303 self.negative_backend.on_records_refreshed(
2304 &self.records,
2305 max_window_tokens,
2306 &|id| self.split_labels.get(id).copied(),
2307 self.ingestion.last_refreshed_sources(),
2308 );
2309 self.epoch_tracker.ensure_loaded()?;
2310 self.refresh_records_by_split_cache()?;
2311 let cached_records_by_split = self
2312 .cached_records_by_split
2313 .as_ref()
2314 .expect("records_by_split cache populated by refresh");
2315 self.epoch_tracker
2316 .reconcile(target_split, cached_records_by_split);
2317 self.ensure_source_state()?;
2318 Ok(())
2319 }
2320
2321 fn ingest_with_weights_fallback(
2326 &mut self,
2327 target_split: SplitLabel,
2328 weights: Option<&HashMap<SourceId, f32>>,
2329 ) -> Result<(), SamplerError> {
2330 match weights {
2331 Some(weights)
2332 if !weights.is_empty()
2333 && !weights
2334 .values()
2335 .all(|&w| w == *weights.values().next().unwrap()) =>
2336 {
2337 self.ingest_internal_with_weights_for_split(target_split, weights)?
2338 }
2339 _ => self.ingest_internal_for_split(target_split)?,
2340 }
2341 Ok(())
2342 }
2343
2344 fn force_ingest_refresh_with_weights_for_split(
2345 &mut self,
2346 target_split: SplitLabel,
2347 weights: &HashMap<SourceId, f32>,
2348 ) -> Result<(), SamplerError> {
2349 if !self.ingestion.has_sources() {
2350 return Ok(());
2351 }
2352 if !self.ingestion_cursors_loaded {
2353 if let Some(state) = self.split_store.load_sampler_state()? {
2354 self.ingestion.load_cursors(&state.source_stream_cursors);
2355 self.ingestion.set_epoch(state.epoch);
2356 self.ingestion.set_epoch_step(state.epoch_step);
2357 }
2358 self.ingestion_cursors_loaded = true;
2359 }
2360 self.ingestion.force_refresh_all_with_weights(weights)?;
2361 self.last_observed_ingest = self.ingestion.total_ingest_count();
2362 self.sync_records_from_cache()?;
2363 let max_window_tokens = self.config.chunking.max_window_tokens;
2364 self.negative_backend.on_records_refreshed(
2365 &self.records,
2366 max_window_tokens,
2367 &|id| self.split_labels.get(id).copied(),
2368 self.ingestion.last_refreshed_sources(),
2369 );
2370 self.epoch_tracker.ensure_loaded()?;
2371 self.refresh_records_by_split_cache()?;
2372 let cached_records_by_split = self
2373 .cached_records_by_split
2374 .as_ref()
2375 .expect("records_by_split cache populated by refresh");
2376 self.epoch_tracker
2377 .reconcile(target_split, cached_records_by_split);
2378 self.ensure_source_state()?;
2379 Ok(())
2380 }
2381
2382 #[allow(clippy::too_many_arguments)]
2394 fn sample_source_triplet_candidate(
2395 &mut self,
2396 source: &str,
2397 target_split: SplitLabel,
2398 recipe_orders: &mut HashMap<RecipeKey, Vec<usize>>,
2399 recipe_positions: &mut HashMap<RecipeKey, usize>,
2400 rng: &mut DeterministicRng,
2401 seen: &mut HashSet<(RecordKey, RecordKey, RecordKey)>,
2402 seen_negatives: &mut IntSet<u64>,
2403 ) -> (Option<(TripletRecipe, SampleTriplet)>, usize) {
2404 let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
2407 if recipes.is_empty() {
2408 return (None, 0);
2409 }
2410 if !recipe_orders.contains_key(source) {
2411 let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2412 let order =
2413 self.recipe_order_weighted_cycled(&recipe_weights, self.triplet_recipe_rr_idx, rng);
2414 recipe_orders.insert(source.to_string(), order);
2415 }
2416 let order = recipe_orders
2417 .get(source)
2418 .expect("recipe order missing for source");
2419 let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2420 let Some(anchor) = self.choose_anchor_record(Some(source), target_split) else {
2421 return (None, 0);
2422 };
2423
2424 let mut attempts = 0usize;
2425 for offset in 0..order.len() {
2426 let idx = order[(*pos + offset) % order.len()];
2427 attempts = attempts.saturating_add(1);
2428 let recipe = recipes[idx].clone();
2429 if let Some(sample) =
2434 self.make_triplet_with_unique_negative(&recipe, &anchor, rng, seen, seen_negatives)
2435 {
2436 *pos = (*pos + offset + 1) % order.len();
2437 return (Some((recipe, sample)), attempts);
2438 }
2439 }
2440
2441 (None, attempts)
2442 }
2443
2444 fn next_pair_batch_inner_with_weights(
2445 &mut self,
2446 target_split: SplitLabel,
2447 weights: Option<&HashMap<SourceId, f32>>,
2448 ) -> Result<SampleBatch, SamplerError> {
2449 self.ingest_with_weights_fallback(target_split, weights)?;
2450 self.ensure_split_has_records(target_split)?;
2451 let sources = self.source_order.clone();
2452 if sources.is_empty() {
2453 if self.triplet_recipes.is_empty() {
2454 return Err(SamplerError::Configuration(
2455 "no triplet recipes available".into(),
2456 ));
2457 }
2458 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2459 let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2460 let recipe_order = self.recipe_order_weighted_cycled(
2461 &recipe_weights,
2462 self.triplet_recipe_rr_idx,
2463 &mut rng,
2464 );
2465 let mut pairs = Vec::new();
2466 let mut seen = HashSet::new();
2467 let mut seen_negatives: IntSet<u64> = IntSet::default();
2471 let mut last_recipe_name = None;
2472 let mut recipe_pos = 0usize;
2473 let mut recipe_steps = 0usize;
2474 let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2475 for _ in 0..attempts {
2476 if pairs.len() >= self.config.batch_size {
2477 break;
2478 }
2479 let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2480 break;
2481 };
2482 let mut triplet = None;
2483 for offset in 0..recipe_order.len() {
2484 let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2485 recipe_steps = recipe_steps.saturating_add(1);
2486 let recipe = self.triplet_recipes[idx].clone();
2487 last_recipe_name = Some(recipe.name.clone());
2488 if let Some(sample) = self.make_triplet_with_unique_negative(
2489 &recipe,
2490 &anchor,
2491 &mut rng,
2492 &mut seen,
2493 &mut seen_negatives,
2494 ) {
2495 triplet = Some((recipe, sample));
2496 recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2497 break;
2498 }
2499 }
2500 if let Some((recipe, triplet)) = triplet {
2501 let SampleTriplet {
2502 recipe: triplet_recipe_name,
2503 anchor,
2504 positive,
2505 negative,
2506 weight,
2507 instruction,
2508 } = triplet;
2509
2510 let record_label = self
2513 .records
2514 .get(&RecordKey::of_chunk(&anchor))
2515 .and_then(|r| r.label.clone());
2516
2517 if pairs.len() < self.config.batch_size {
2518 pairs.push(SamplePair {
2519 recipe: triplet_recipe_name.clone(),
2520 anchor: anchor.clone(),
2521 positive: positive.clone(),
2522 weight,
2523 instruction: instruction.clone(),
2524 label: record_label.clone().unwrap_or(PairLabel::Positive),
2525 reason: None,
2526 });
2527 }
2528 if pairs.len() < self.config.batch_size {
2529 pairs.push(SamplePair {
2530 recipe: triplet_recipe_name,
2531 anchor,
2532 positive: negative,
2533 weight,
2534 instruction,
2535 label: record_label.unwrap_or(PairLabel::Negative),
2536 reason: Some(strategy_reason(&recipe.negative_strategy).to_string()),
2537 });
2538 }
2539 }
2540 }
2541 if recipe_steps > 0 {
2542 self.triplet_recipe_rr_idx =
2543 self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2544 }
2545 self.rng = rng;
2546 pad_with_reuse(&mut pairs, self.config.batch_size);
2547 if pairs.len() == self.config.batch_size {
2548 return Ok(SampleBatch { pairs });
2549 }
2550 return Err(SamplerError::Exhausted(
2551 last_recipe_name
2552 .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2553 .to_string(),
2554 ));
2555 }
2556
2557 let mut pairs = Vec::new();
2558 let mut seen = HashSet::new();
2559 let mut seen_negatives: IntSet<u64> = IntSet::default();
2561 let mut source_steps = 0usize;
2562 let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2563 let mut source_idx = self.source_cycle_idx % sources.len();
2564 let mut cycle_sources = self.shuffled_source_cycle(cycle);
2565 let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2568 Some(w) if Self::weights_are_selective(w) => Some(w),
2569 _ => None,
2570 };
2571 let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2572 let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2573 let mut recipe_steps = 0usize;
2574 let max_recipe_len = sources
2575 .iter()
2576 .map(|source| self.triplet_recipe_count_for_source(source))
2577 .max()
2578 .unwrap_or(1)
2579 .max(1);
2580 let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2581 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2582 for _ in 0..attempts {
2583 if pairs.len() >= self.config.batch_size {
2584 break;
2585 }
2586 let weighted_source;
2587 let source = if let Some(w) = weighted {
2588 match self.select_weighted_source(&sources, w) {
2589 Some(s) => {
2590 weighted_source = s;
2591 weighted_source.as_str()
2592 }
2593 None => break,
2594 }
2595 } else {
2596 cycle_sources[source_idx].as_str()
2597 };
2598 let (triplet, attempts_used) = self.sample_source_triplet_candidate(
2599 source,
2600 target_split,
2601 &mut recipe_orders,
2602 &mut recipe_positions,
2603 &mut rng,
2604 &mut seen,
2605 &mut seen_negatives,
2606 );
2607 recipe_steps = recipe_steps.saturating_add(attempts_used);
2608 if let Some((recipe, triplet)) = triplet {
2609 let SampleTriplet {
2610 recipe: triplet_recipe_name,
2611 anchor,
2612 positive,
2613 negative,
2614 weight,
2615 instruction,
2616 } = triplet;
2617
2618 let record_label = self
2620 .records
2621 .get(&RecordKey::of_chunk(&anchor))
2622 .and_then(|r| r.label.clone());
2623
2624 if pairs.len() < self.config.batch_size {
2625 pairs.push(SamplePair {
2626 recipe: triplet_recipe_name.clone(),
2627 anchor: anchor.clone(),
2628 positive: positive.clone(),
2629 weight,
2630 instruction: instruction.clone(),
2631 label: record_label.clone().unwrap_or(PairLabel::Positive),
2632 reason: None,
2633 });
2634 }
2635 if pairs.len() < self.config.batch_size {
2636 pairs.push(SamplePair {
2637 recipe: triplet_recipe_name,
2638 anchor,
2639 positive: negative,
2640 weight,
2641 instruction,
2642 label: record_label.unwrap_or(PairLabel::Negative),
2643 reason: Some(strategy_reason(&recipe.negative_strategy).to_string()),
2644 });
2645 }
2646 }
2647 source_idx += 1;
2648 source_steps += 1;
2649 if source_idx >= cycle_sources.len() {
2650 source_idx = 0;
2651 cycle = cycle.saturating_add(1);
2652 cycle_sources = self.shuffled_source_cycle(cycle);
2653 }
2654 }
2655 if recipe_steps > 0 {
2656 self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2657 }
2658 self.rng = rng;
2659 pad_with_reuse(&mut pairs, self.config.batch_size);
2660 if pairs.len() == self.config.batch_size {
2661 self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2662 self.source_state_dirty = sources.len() > 1;
2663 return Ok(SampleBatch { pairs });
2664 }
2665 Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
2666 }
2667
2668 fn try_emit_text_sample(
2673 &mut self,
2674 sample: &TextSample,
2675 batch_texts: &mut HashSet<String>,
2676 ) -> bool {
2677 batch_texts.insert(sample.chunk.text.clone())
2678 && self
2679 .emitted_text_hashes
2680 .entry(RecordKey::of_chunk(&sample.chunk))
2681 .or_default()
2682 .insert(stable_hash_str(0, &sample.chunk.text))
2683 }
2684
2685 fn next_text_batch_inner_with_weights(
2686 &mut self,
2687 target_split: SplitLabel,
2688 weights: Option<&HashMap<SourceId, f32>>,
2689 ) -> Result<TextBatch, SamplerError> {
2690 self.ingest_with_weights_fallback(target_split, weights)?;
2691 self.ensure_split_has_records(target_split)?;
2692 let sources = self.source_order.clone();
2693 if sources.is_empty() {
2694 if self.text_recipes.is_empty() {
2695 return Err(SamplerError::Configuration(
2696 "no text recipes configured".into(),
2697 ));
2698 }
2699 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2700 let recipe_weights: Vec<f32> = self.text_recipes.iter().map(|r| r.weight).collect();
2701 let recipe_order = self.text_recipe_order_weighted_cycled(
2702 &recipe_weights,
2703 self.text_recipe_rr_idx,
2704 &mut rng,
2705 );
2706 let mut samples = Vec::new();
2707 let mut batch_texts = HashSet::new();
2708 let mut last_recipe_name = None;
2709 let mut recipe_pos = 0usize;
2710 let mut recipe_steps = 0usize;
2711 let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2712 for _ in 0..attempts {
2713 if samples.len() >= self.config.batch_size {
2714 break;
2715 }
2716 let recipe_idx = recipe_order[recipe_pos];
2717 recipe_pos = (recipe_pos + 1) % recipe_order.len();
2718 recipe_steps = recipe_steps.saturating_add(1);
2719 let recipe = self.text_recipes[recipe_idx].clone();
2720 last_recipe_name = Some(recipe.name.clone());
2721 if let Some(sample) =
2722 self.make_text_sample_for_split(&recipe, None, target_split, &mut rng)
2723 && self.try_emit_text_sample(&sample, &mut batch_texts)
2724 {
2725 samples.push(sample);
2726 }
2727 }
2728 if recipe_steps > 0 {
2729 self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2730 }
2731 self.rng = rng;
2732 pad_with_reuse(&mut samples, self.config.batch_size);
2733 if samples.len() == self.config.batch_size {
2734 return Ok(TextBatch { samples });
2735 }
2736 return Err(SamplerError::Exhausted(
2737 last_recipe_name
2738 .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TEXT))
2739 .to_string(),
2740 ));
2741 }
2742
2743 let mut samples = Vec::new();
2744 let mut batch_texts = HashSet::new();
2745 let mut source_steps = 0usize;
2746 let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2747 let mut idx = self.source_cycle_idx % sources.len();
2748 let mut cycle_sources = self.shuffled_source_cycle(cycle);
2749 let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2750 let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2751 let mut recipe_steps = 0usize;
2752 let max_recipe_len = sources
2753 .iter()
2754 .map(|source| self.text_recipes_for_source(source).len())
2755 .max()
2756 .unwrap_or(1)
2757 .max(1);
2758 let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2759 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2760 let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2763 Some(w) if Self::weights_are_selective(w) => Some(w),
2764 _ => None,
2765 };
2766 for _ in 0..attempts {
2767 if samples.len() >= self.config.batch_size {
2768 break;
2769 }
2770 let weighted_source;
2771 let source = if let Some(w) = weighted {
2772 match self.select_weighted_source(&sources, w) {
2773 Some(s) => {
2774 weighted_source = s;
2775 weighted_source.as_str()
2776 }
2777 None => break,
2778 }
2779 } else {
2780 cycle_sources[idx].as_str()
2781 };
2782 let recipes = self.text_recipes_for_source(source).to_vec();
2783 if recipes.is_empty() {
2784 idx += 1;
2785 source_steps += 1;
2786 if idx >= cycle_sources.len() {
2787 idx = 0;
2788 cycle = cycle.saturating_add(1);
2789 cycle_sources = self.shuffled_source_cycle(cycle);
2790 }
2791 continue;
2792 }
2793 if !recipe_orders.contains_key(source) {
2794 let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2795 let order = self.text_recipe_order_weighted_cycled(
2796 &recipe_weights,
2797 self.text_recipe_rr_idx,
2798 &mut rng,
2799 );
2800 recipe_orders.insert(source.to_string(), order);
2801 }
2802 let order = recipe_orders
2803 .get(source)
2804 .expect("recipe order missing for source");
2805 let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2806 let mut sample: Option<(TextRecipe, TextSample)> = None;
2807 for offset in 0..order.len() {
2808 let recipe_idx = order[(*pos + offset) % order.len()];
2809 let recipe = recipes[recipe_idx].clone();
2810 if let Some(item) =
2811 self.make_text_sample_for_split(&recipe, Some(source), target_split, &mut rng)
2812 && self.try_emit_text_sample(&item, &mut batch_texts)
2813 {
2814 recipe_steps = recipe_steps.saturating_add(offset + 1);
2815 *pos = (*pos + offset + 1) % order.len();
2816 sample = Some((recipe, item));
2817 break;
2818 }
2819 }
2820 if sample.is_none() {
2821 recipe_steps = recipe_steps.saturating_add(order.len());
2822 }
2823 if let Some((_recipe, sample)) = sample {
2824 samples.push(sample);
2825 }
2826 idx += 1;
2827 source_steps += 1;
2828 if idx >= cycle_sources.len() {
2829 idx = 0;
2830 cycle = cycle.saturating_add(1);
2831 cycle_sources = self.shuffled_source_cycle(cycle);
2832 }
2833 }
2834 if samples.len() != self.config.batch_size {
2835 pad_with_reuse(&mut samples, self.config.batch_size);
2836 }
2837 if samples.len() != self.config.batch_size {
2838 self.rng = rng;
2839 return Err(SamplerError::Exhausted(RECIPE_LABEL_TEXT.into()));
2840 }
2841 self.rng = rng;
2842 self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2843 self.source_state_dirty = sources.len() > 1;
2844 if recipe_steps > 0 {
2845 self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2846 }
2847 Ok(TextBatch { samples })
2848 }
2849
2850 fn next_triplet_batch_inner_with_weights(
2851 &mut self,
2852 target_split: SplitLabel,
2853 weights: Option<&HashMap<SourceId, f32>>,
2854 ) -> Result<TripletBatch, SamplerError> {
2855 self.ingest_with_weights_fallback(target_split, weights)?;
2856 self.ensure_split_has_records(target_split)?;
2857 let sources = self.source_order.clone();
2858 if sources.is_empty() {
2859 if self.triplet_recipes.is_empty() {
2860 return Err(SamplerError::Configuration(
2861 "no triplet recipes configured".into(),
2862 ));
2863 }
2864 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2865 let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2866 let recipe_order = self.recipe_order_weighted_cycled(
2867 &recipe_weights,
2868 self.triplet_recipe_rr_idx,
2869 &mut rng,
2870 );
2871 let mut triplets = Vec::new();
2872 let mut seen = HashSet::new();
2873 let mut seen_negatives: IntSet<u64> = IntSet::default();
2875 let mut last_recipe_name = None;
2876 let mut recipe_pos = 0usize;
2877 let mut recipe_steps = 0usize;
2878 let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2879 for _ in 0..attempts {
2880 if triplets.len() >= self.config.batch_size {
2881 break;
2882 }
2883 let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2884 break;
2885 };
2886 let mut triplet = None;
2887 for offset in 0..recipe_order.len() {
2888 let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2889 recipe_steps = recipe_steps.saturating_add(1);
2890 let recipe = self.triplet_recipes[idx].clone();
2891 last_recipe_name = Some(recipe.name.clone());
2892 if let Some(sample) = self.make_triplet_with_unique_negative(
2893 &recipe,
2894 &anchor,
2895 &mut rng,
2896 &mut seen,
2897 &mut seen_negatives,
2898 ) {
2899 triplet = Some(sample);
2900 recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2901 break;
2902 }
2903 }
2904 if let Some(triplet) = triplet {
2905 triplets.push(triplet);
2906 }
2907 }
2908 if recipe_steps > 0 {
2909 self.triplet_recipe_rr_idx =
2910 self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2911 }
2912 self.rng = rng;
2913 pad_with_reuse(&mut triplets, self.config.batch_size);
2914 if triplets.len() == self.config.batch_size {
2915 return Ok(TripletBatch { triplets });
2916 }
2917 return Err(SamplerError::Exhausted(
2918 last_recipe_name
2919 .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2920 .to_string(),
2921 ));
2922 }
2923
2924 let mut triplets: Vec<SampleTriplet> = Vec::new();
2926 let mut seen: HashSet<(RecordKey, RecordKey, RecordKey)> = HashSet::new();
2927 let mut seen_negatives: IntSet<u64> = IntSet::default();
2929 let mut source_steps = 0usize;
2930 let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2931 let mut source_idx = self.source_cycle_idx % sources.len();
2932 let mut cycle_sources = self.shuffled_source_cycle(cycle);
2933 let weighted: Option<&HashMap<SourceId, f32>> = match weights {
2936 Some(w) if Self::weights_are_selective(w) => Some(w),
2937 _ => None,
2938 };
2939 let mut recipe_steps = 0usize;
2940 let max_recipe_len = sources
2941 .iter()
2942 .map(|source| self.triplet_recipe_count_for_source(source.as_str()))
2943 .max()
2944 .unwrap_or(1)
2945 .max(1);
2946 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2947
2948 struct SlotPlan {
2953 anchor: Arc<DataRecord>,
2954 recipes: Vec<TripletRecipe>,
2955 fork_seed: u64,
2956 }
2957 let target_slots = self.config.batch_size * 4 * max_recipe_len;
2958 let mut slot_plans: Vec<SlotPlan> = Vec::with_capacity(target_slots);
2959
2960 for _ in 0..target_slots {
2961 if slot_plans.len() >= target_slots {
2962 break;
2963 }
2964 let weighted_source;
2965 let source = if let Some(w) = weighted {
2966 match self.select_weighted_source(&sources, w) {
2967 Some(s) => {
2968 weighted_source = s;
2969 weighted_source.as_str()
2970 }
2971 None => break,
2972 }
2973 } else {
2974 cycle_sources[source_idx].as_str()
2975 };
2976 let (recipes, _) = self.resolve_source_triplet_plan(source);
2977 if !recipes.is_empty() {
2978 let fork_seed = rng.next_u64();
2979 if let Some(anchor) = self.choose_anchor_record(Some(source), target_split) {
2980 slot_plans.push(SlotPlan {
2981 anchor,
2982 recipes,
2983 fork_seed,
2984 });
2985 recipe_steps = recipe_steps.saturating_add(1);
2986 }
2987 }
2988 source_idx += 1;
2989 source_steps += 1;
2990 if source_idx >= cycle_sources.len() {
2991 source_idx = 0;
2992 cycle = cycle.saturating_add(1);
2993 cycle_sources = self.shuffled_source_cycle(cycle);
2994 }
2995 }
2996 self.rng = rng;
2997
2998 struct SlotCandidate {
3003 recipe: TripletRecipe,
3004 anchor: Arc<DataRecord>,
3005 anchor_chunk: RecordChunk,
3006 positive_chunk: RecordChunk,
3007 anchor_raw_text: String,
3008 }
3009 let mut raw_candidates: Vec<(usize, Option<SlotCandidate>)> = slot_plans
3010 .par_iter()
3011 .enumerate()
3012 .map(|(slot_idx, plan)| {
3013 let mut fork_rng = DeterministicRng::new(plan.fork_seed);
3014 let weights: Vec<f32> = plan.recipes.iter().map(|r| r.weight).collect();
3017 let order = weighted_recipe_order(&weights, &mut fork_rng);
3018 let mut candidate = None;
3019 for &idx in &order {
3020 let recipe = &plan.recipes[idx];
3021 if let Some((ac, pc, raw)) =
3022 self.select_anchor_positive_for_recipe(recipe, &plan.anchor, &mut fork_rng)
3023 {
3024 candidate = Some(SlotCandidate {
3025 recipe: recipe.clone(),
3026 anchor: Arc::clone(&plan.anchor),
3027 anchor_chunk: ac,
3028 positive_chunk: pc,
3029 anchor_raw_text: raw,
3030 });
3031 break;
3032 }
3033 }
3034 (slot_idx, candidate)
3035 })
3036 .collect();
3037
3038 raw_candidates.sort_unstable_by_key(|(i, _)| *i);
3043 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3044 for (_, candidate) in raw_candidates {
3045 if triplets.len() >= self.config.batch_size {
3046 break;
3047 }
3048 let Some(sc) = candidate else { continue };
3049 let (final_anchor, final_positive) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0
3052 {
3053 (sc.positive_chunk, sc.anchor_chunk)
3054 } else {
3055 (sc.anchor_chunk, sc.positive_chunk)
3056 };
3057 if !sc.recipe.allow_same_anchor_positive && final_anchor.text == final_positive.text {
3058 continue;
3059 }
3060 let mut accepted: Option<(RecordChunk, f32, String)> = None;
3063 for _ in 0..NEG_DEDUP_RETRIES {
3064 let (negative_record, fallback_used) = match self.select_negative_record(
3065 &sc.anchor,
3066 &sc.recipe.negative_strategy,
3067 Some(&sc.anchor_raw_text),
3068 &mut rng,
3069 ) {
3070 Some(pair) => pair,
3071 None => break,
3072 };
3073 let mut negative_chunk = match self.select_chunk_parallel(
3074 &negative_record,
3075 &sc.recipe.negative_selector,
3076 &mut rng,
3077 ) {
3078 Some(c) => c,
3079 None => continue,
3080 };
3081 self.decorate_chunk_parallel(&negative_record, &mut negative_chunk, &mut rng);
3082 if negative_chunk.text == final_positive.text
3083 || negative_chunk.text == final_anchor.text
3084 {
3085 continue;
3086 }
3087 let key = (
3088 RecordKey::of_chunk(&final_anchor),
3089 RecordKey::of_chunk(&final_positive),
3090 RecordKey::of_chunk(&negative_chunk),
3091 );
3092 let neg_h = negative_text_hash(&negative_chunk);
3093 if seen.contains(&key) || seen_negatives.contains(&neg_h) {
3095 continue;
3096 }
3097 let chunk_weight = self.triplet_chunk_weight(
3098 &sc.recipe,
3099 &final_anchor,
3100 &final_positive,
3101 &negative_chunk,
3102 );
3103 let weight = sc.recipe.weight * chunk_weight;
3104 let recipe_name = if fallback_used {
3105 format!("{}_fallback_same_split", sc.recipe.name)
3106 } else {
3107 sc.recipe.name.to_string()
3108 };
3109 seen.insert(key);
3110 seen_negatives.insert(neg_h);
3111 accepted = Some((negative_chunk, weight, recipe_name));
3112 break;
3113 }
3114 let Some((negative_chunk, weight, recipe_name)) = accepted else {
3115 continue;
3116 };
3117 let triplet = SampleTriplet {
3118 recipe: recipe_name,
3119 anchor: final_anchor,
3120 positive: final_positive,
3121 negative: negative_chunk,
3122 weight,
3123 instruction: sc.recipe.instruction.as_ref().map(|s| s.to_string()),
3124 };
3125 if triplets.len() < self.config.batch_size {
3126 triplets.push(triplet);
3127 }
3128 }
3129 self.rng = rng;
3130
3131 if recipe_steps > 0 {
3132 self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
3133 }
3134 pad_with_reuse(&mut triplets, self.config.batch_size);
3135 if triplets.len() == self.config.batch_size {
3136 self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
3137 self.source_state_dirty = sources.len() > 1;
3138 let batch = TripletBatch { triplets };
3139 return Ok(batch);
3140 }
3141 Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
3142 }
3143
3144 #[cfg(test)]
3150 #[cfg(feature = "bm25-mining")]
3151 fn bm25_backend_mut(&mut self) -> &mut backends::Bm25Backend {
3152 self.negative_backend
3153 .as_any_mut()
3154 .downcast_mut::<backends::Bm25Backend>()
3155 .expect("bm25_backend_mut: negative_backend is Bm25Backend when bm25-mining feature is active")
3156 }
3157
3158 #[cfg(test)]
3160 fn recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
3161 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3162 let result = self.recipe_order_weighted_shuffled(weights, &mut rng);
3163 self.rng = rng;
3164 result
3165 }
3166
3167 #[cfg(test)]
3169 fn recipe_order_weighted_cycled_seeded(
3170 &mut self,
3171 weights: &[f32],
3172 rr_idx: usize,
3173 ) -> Vec<usize> {
3174 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3175 let result = self.recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
3176 self.rng = rng;
3177 result
3178 }
3179
3180 #[cfg(test)]
3182 fn text_recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
3183 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3184 let result = self.text_recipe_order_weighted_shuffled(weights, &mut rng);
3185 self.rng = rng;
3186 result
3187 }
3188
3189 #[cfg(test)]
3191 fn text_recipe_order_weighted_cycled_seeded(
3192 &mut self,
3193 weights: &[f32],
3194 rr_idx: usize,
3195 ) -> Vec<usize> {
3196 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3197 let result = self.text_recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
3198 self.rng = rng;
3199 result
3200 }
3201
3202 #[cfg(test)]
3204 fn select_negative_record_seeded(
3205 &mut self,
3206 anchor_record: &DataRecord,
3207 strategy: &NegativeStrategy,
3208 anchor_query_text: Option<&str>,
3209 ) -> Option<(Arc<DataRecord>, bool)> {
3210 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3211 let result =
3212 self.select_negative_record(anchor_record, strategy, anchor_query_text, &mut rng);
3213 self.rng = rng;
3214 result
3215 }
3216
3217 #[cfg(test)]
3220 fn make_triplet_with_unique_negative_seeded(
3221 &mut self,
3222 recipe: &TripletRecipe,
3223 anchor: &DataRecord,
3224 ) -> Option<SampleTriplet> {
3225 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3226 let mut seen = HashSet::new();
3227 let mut seen_negatives: IntSet<u64> = IntSet::default();
3228 let result = self.make_triplet_with_unique_negative(
3229 recipe,
3230 anchor,
3231 &mut rng,
3232 &mut seen,
3233 &mut seen_negatives,
3234 );
3235 self.rng = rng;
3236 result
3237 }
3238
3239 #[cfg(test)]
3241 fn decorate_chunk_seeded(&mut self, record: &DataRecord, chunk: &mut RecordChunk) {
3242 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
3243 self.decorate_chunk(record, chunk, &mut rng);
3244 self.rng = rng;
3245 }
3246
3247 #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3253 fn bm25_fallback_stats(&self) -> (u64, u64) {
3254 self.negative_backend.bm25_fallback_stats()
3255 }
3256
3257 #[cfg(test)]
3262 #[cfg(feature = "bm25-mining")]
3263 fn bm25_ranked_candidates(&mut self, anchor: &crate::data::DataRecord) -> Vec<RecordKey> {
3264 let split = self
3265 .split_store
3266 .label_for(&RecordKey::from(anchor))
3267 .unwrap_or(SplitLabel::Train);
3268 self.negative_backend
3269 .as_any_mut()
3270 .downcast_mut::<backends::Bm25Backend>()
3271 .expect("bm25_ranked_candidates: Bm25Backend")
3272 .ranked_candidates_pub(anchor, split)
3273 }
3274}
3275
3276fn weighted_recipe_order(weights: &[f32], rng: &mut DeterministicRng) -> Vec<usize> {
3287 let nonzero: Vec<(usize, f32)> = weights
3288 .iter()
3289 .enumerate()
3290 .filter(|(_, w)| **w > 0.0)
3291 .map(|(i, &w)| (i, w))
3292 .collect();
3293 if nonzero.is_empty() {
3294 return Vec::new();
3295 }
3296 let w_min = nonzero
3297 .iter()
3298 .map(|(_, w)| *w)
3299 .fold(f32::INFINITY, f32::min);
3300 let mut order: Vec<usize> = Vec::new();
3301 for (recipe_idx, w) in &nonzero {
3302 let tickets = ((w / w_min).round() as usize).clamp(1, RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER);
3303 for _ in 0..tickets {
3304 order.push(*recipe_idx);
3305 }
3306 }
3307 order.shuffle(rng);
3308 order
3309}
3310
3311fn same_selector_pair_is_valid(
3312 anchor_chunk: &RecordChunk,
3313 positive_chunk: &RecordChunk,
3314 enforce_window_pair: bool,
3315) -> bool {
3316 if triplet_chunk_key(anchor_chunk) == triplet_chunk_key(positive_chunk) {
3317 return false;
3318 }
3319 if !enforce_window_pair {
3320 return true;
3321 }
3322 matches!(
3323 (&anchor_chunk.view, &positive_chunk.view),
3324 (ChunkView::Window { .. }, ChunkView::Window { .. })
3325 )
3326}
3327
3328impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSampler<S> {
3329 pub fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
3331 let inner = TripletSamplerInner::new(config, split_store);
3332 Self {
3333 inner: Mutex::new(inner),
3334 }
3335 }
3336
3337 pub fn new_with_chunker(
3339 config: SamplerConfig,
3340 split_store: Arc<S>,
3341 chunker: Arc<dyn ChunkingAlgorithm>,
3342 ) -> Self {
3343 let inner = TripletSamplerInner::new_with_chunker(config, split_store, chunker);
3344 Self {
3345 inner: Mutex::new(inner),
3346 }
3347 }
3348
3349 #[deprecated(
3355 note = "samples all sources uniformly; call next_pair_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3356 )]
3357 pub fn next_pair_batch_for_split(
3358 &self,
3359 split: SplitLabel,
3360 ) -> Result<SampleBatch, SamplerError> {
3361 self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3362 }
3363
3364 #[deprecated(
3370 note = "samples all sources uniformly; call next_text_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3371 )]
3372 pub fn next_text_batch_for_split(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3373 self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3374 }
3375
3376 #[deprecated(
3382 note = "samples all sources uniformly; call next_triplet_batch_with_weights_for_split(split, &weights) with an explicit per-source weight map to honor a data mixture"
3383 )]
3384 pub fn next_triplet_batch_for_split(
3385 &self,
3386 split: SplitLabel,
3387 ) -> Result<TripletBatch, SamplerError> {
3388 self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3389 }
3390
3391 fn next_batch_with_retry<B>(
3395 &self,
3396 split: SplitLabel,
3397 weights: &HashMap<SourceId, f32>,
3398 inner_fn: impl Fn(
3399 &mut TripletSamplerInner<S>,
3400 SplitLabel,
3401 Option<&HashMap<SourceId, f32>>,
3402 ) -> Result<B, SamplerError>,
3403 ) -> Result<B, SamplerError> {
3404 let mut inner = self.inner.lock().unwrap();
3405 inner.ensure_split_allowed(split)?;
3406 if !weights.is_empty() {
3409 inner.ingestion.validate_weights(weights)?;
3410 }
3411 inner.ensure_ingestion_cursors_loaded()?;
3414 inner.ingestion.increment_epoch_step();
3416 for attempt in 0..=EXHAUSTION_RETRY_LIMIT {
3417 match inner_fn(&mut inner, split, Some(weights)) {
3418 Ok(batch) => return Ok(batch),
3419 Err(SamplerError::Exhausted(_)) if attempt < EXHAUSTION_RETRY_LIMIT => {
3420 inner.force_ingest_refresh_with_weights_for_split(split, weights)?;
3421 }
3422 Err(err) => return Err(err),
3423 }
3424 }
3425 unreachable!()
3427 }
3428
3429 pub fn next_pair_batch_with_weights_for_split(
3437 &self,
3438 split: SplitLabel,
3439 weights: &HashMap<SourceId, f32>,
3440 ) -> Result<SampleBatch, SamplerError> {
3441 self.next_batch_with_retry(split, weights, |inner, split, weights| {
3442 inner.next_pair_batch_inner_with_weights(split, weights)
3443 })
3444 }
3445
3446 pub fn next_text_batch_with_weights_for_split(
3453 &self,
3454 split: SplitLabel,
3455 weights: &HashMap<SourceId, f32>,
3456 ) -> Result<TextBatch, SamplerError> {
3457 self.next_batch_with_retry(split, weights, |inner, split, weights| {
3458 inner.next_text_batch_inner_with_weights(split, weights)
3459 })
3460 }
3461
3462 pub fn next_triplet_batch_with_weights_for_split(
3469 &self,
3470 split: SplitLabel,
3471 weights: &HashMap<SourceId, f32>,
3472 ) -> Result<TripletBatch, SamplerError> {
3473 self.next_batch_with_retry(split, weights, |inner, split, weights| {
3474 inner.next_triplet_batch_inner_with_weights(split, weights)
3475 })
3476 }
3477
3478 #[deprecated(
3484 note = "samples all sources uniformly; call prefetch_triplet_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3485 )]
3486 pub fn prefetch_triplet_batches(
3487 self: Arc<Self>,
3488 split: SplitLabel,
3489 capacity: usize,
3490 ) -> BatchPrefetcher<TripletBatch> {
3491 BatchPrefetcher::new(capacity, move || {
3492 self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3493 })
3494 }
3495
3496 pub fn prefetch_triplet_batches_with_weights(
3498 self: Arc<Self>,
3499 split: SplitLabel,
3500 capacity: usize,
3501 weights: HashMap<SourceId, f32>,
3502 ) -> BatchPrefetcher<TripletBatch> {
3503 BatchPrefetcher::new(capacity, move || {
3504 self.next_triplet_batch_with_weights_for_split(split, &weights)
3505 })
3506 }
3507
3508 #[deprecated(
3514 note = "samples all sources uniformly; call prefetch_pair_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3515 )]
3516 pub fn prefetch_pair_batches(
3517 self: Arc<Self>,
3518 split: SplitLabel,
3519 capacity: usize,
3520 ) -> BatchPrefetcher<SampleBatch> {
3521 BatchPrefetcher::new(capacity, move || {
3522 self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3523 })
3524 }
3525
3526 pub fn prefetch_pair_batches_with_weights(
3528 self: Arc<Self>,
3529 split: SplitLabel,
3530 capacity: usize,
3531 weights: HashMap<SourceId, f32>,
3532 ) -> BatchPrefetcher<SampleBatch> {
3533 BatchPrefetcher::new(capacity, move || {
3534 self.next_pair_batch_with_weights_for_split(split, &weights)
3535 })
3536 }
3537
3538 #[deprecated(
3544 note = "samples all sources uniformly; call prefetch_text_batches_with_weights(split, capacity, weights) with an explicit per-source weight map to honor a data mixture"
3545 )]
3546 pub fn prefetch_text_batches(
3547 self: Arc<Self>,
3548 split: SplitLabel,
3549 capacity: usize,
3550 ) -> BatchPrefetcher<TextBatch> {
3551 BatchPrefetcher::new(capacity, move || {
3552 self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3553 })
3554 }
3555
3556 pub fn prefetch_text_batches_with_weights(
3558 self: Arc<Self>,
3559 split: SplitLabel,
3560 capacity: usize,
3561 weights: HashMap<SourceId, f32>,
3562 ) -> BatchPrefetcher<TextBatch> {
3563 BatchPrefetcher::new(capacity, move || {
3564 self.next_text_batch_with_weights_for_split(split, &weights)
3565 })
3566 }
3567
3568 pub fn text_recipes(&self) -> Vec<TextRecipe> {
3570 let inner = self.inner.lock().unwrap();
3571 inner.text_recipes().to_vec()
3572 }
3573
3574 pub fn register_source(
3579 &self,
3580 source: Box<dyn DataSource + 'static>,
3581 ) -> Result<(), SamplerError> {
3582 let mut inner = self.inner.lock().unwrap();
3583 inner.register_source(source)
3584 }
3585
3586 pub fn set_epoch(&self, epoch: u64) -> Result<(), SamplerError> {
3588 let mut inner = self.inner.lock().unwrap();
3589 inner.set_epoch(epoch)
3590 }
3591
3592 pub fn save_sampler_state(&self, save_to: Option<&Path>) -> Result<(), SamplerError> {
3597 let mut inner = self.inner.lock().unwrap();
3598 inner.save_sampler_state(save_to)
3599 }
3600
3601 #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3609 pub fn bm25_fallback_stats(&self) -> (u64, u64) {
3610 let inner = self.inner.lock().unwrap();
3611 inner.bm25_fallback_stats()
3612 }
3613}
3614
3615impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> Sampler for TripletSampler<S> {
3616 fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
3617 self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3618 }
3619
3620 fn next_pair_batch_with_weights(
3621 &self,
3622 split: SplitLabel,
3623 weights: &HashMap<SourceId, f32>,
3624 ) -> Result<SampleBatch, SamplerError> {
3625 self.next_pair_batch_with_weights_for_split(split, weights)
3626 }
3627
3628 fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3629 self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3630 }
3631
3632 fn next_text_batch_with_weights(
3633 &self,
3634 split: SplitLabel,
3635 weights: &HashMap<SourceId, f32>,
3636 ) -> Result<TextBatch, SamplerError> {
3637 self.next_text_batch_with_weights_for_split(split, weights)
3638 }
3639
3640 fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
3641 self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3642 }
3643
3644 fn next_triplet_batch_with_weights(
3645 &self,
3646 split: SplitLabel,
3647 weights: &HashMap<SourceId, f32>,
3648 ) -> Result<TripletBatch, SamplerError> {
3649 self.next_triplet_batch_with_weights_for_split(split, weights)
3650 }
3651
3652 fn save_sampler_state(&self, save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
3653 TripletSampler::save_sampler_state(self, save_to)
3654 }
3655}
3656
3657fn roles_match(target: &SectionRole, candidate: &SectionRole) -> bool {
3658 target == candidate
3659}
3660
3661fn role_cursor_key(record_key: &RecordKey, role: &SectionRole) -> (RecordKey, String) {
3662 (record_key.clone(), role_label(role))
3663}
3664
3665fn role_label(role: &SectionRole) -> String {
3666 match role {
3667 SectionRole::Anchor => ROLE_LABEL_ANCHOR.into(),
3668 SectionRole::Context => ROLE_LABEL_CONTEXT.into(),
3669 }
3670}
3671
3672fn taxonomy_value(record: &DataRecord, field: MetadataKey) -> Option<&str> {
3673 record.taxonomy.iter().find_map(|entry| field.strip(entry))
3674}
3675
3676fn strategy_reason(strategy: &NegativeStrategy) -> &'static str {
3677 match strategy {
3678 NegativeStrategy::WrongPublicationDate => NEG_REASON_WRONG_DATE,
3679 NegativeStrategy::WrongArticle => NEG_REASON_WRONG_ARTICLE,
3680 NegativeStrategy::QuestionAnswerMismatch => NEG_REASON_WRONG_QA,
3681 NegativeStrategy::SameRecord => NEG_REASON_SAME_RECORD,
3682 }
3683}
3684
3685#[cfg(test)]
3693fn text_dedup_key(chunk: &RecordChunk) -> (String, String) {
3694 (chunk.record_id.clone(), chunk.text.clone())
3695}
3696
3697fn triplet_chunk_key(chunk: &RecordChunk) -> String {
3698 match &chunk.view {
3699 ChunkView::Window { index, .. } => {
3700 format!("{}|{}|w|{}", chunk.record_id, chunk.section_idx, index)
3701 }
3702 ChunkView::SummaryFallback { strategy, .. } => {
3703 format!("{}|{}|s|{}", chunk.record_id, chunk.section_idx, strategy)
3704 }
3705 }
3706}
3707
3708const NEG_DEDUP_RETRIES: usize = 4;
3715
3716#[inline]
3722fn negative_text_hash(chunk: &RecordChunk) -> u64 {
3723 stable_hash_str(0, &chunk.text)
3724}
3725
3726fn pad_with_reuse<T: Clone>(items: &mut Vec<T>, target: usize) {
3727 if items.is_empty() || items.len() >= target {
3728 return;
3729 }
3730 let seed = items.clone();
3731 let base_len = seed.len();
3732 for idx in 0..(target - items.len()) {
3733 items.push(seed[idx % base_len].clone());
3734 }
3735}
3736
3737#[cfg(test)]
3738mod tests;