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