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_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, RecordId, 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<RecordId, Arc<DataRecord>>,
293 rng: DeterministicRng,
295 triplet_recipes: Vec<TripletRecipe>,
297 text_recipes: Vec<TextRecipe>,
299 source_triplet_recipes: HashMap<SourceId, Vec<TripletRecipe>>,
301 sources_with_long_sections: HashSet<SourceId>,
303 source_text_recipes: HashMap<SourceId, Vec<TextRecipe>>,
305 using_config_triplet_recipes: bool,
307 using_config_text_recipes: bool,
309 last_observed_ingest: u64,
311 epoch_tracker: EpochTracker,
313 chunk_cursors: HashMap<(RecordId, usize), usize>,
315 role_cursors: HashMap<(RecordId, String), usize>,
317 negative_backend: Box<dyn backends::NegativeBackend>,
319 chunk_index: HashMap<RecordId, RecordId>,
321 source_order: Vec<SourceId>,
323 source_cycle_idx: usize,
325 source_state_loaded: bool,
327 ingestion_cursors_loaded: bool,
329 source_state_dirty: bool,
331 source_record_indices: HashMap<SourceId, Vec<usize>>,
333 source_record_cursors: HashMap<SourceId, usize>,
335 triplet_recipe_rr_idx: usize,
337 emitted_text_hashes: HashMap<RecordId, HashSet<u64>>,
346 text_recipe_rr_idx: usize,
348 epoch: u64,
350
351 source_wrapped: HashMap<SourceId, bool>,
353}
354
355impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSamplerInner<S> {
356 fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
357 Self::new_with_chunker(config, split_store, Arc::new(SlidingWindowChunker))
358 }
359
360 fn new_with_chunker(
361 config: SamplerConfig,
362 split_store: Arc<S>,
363 chunker: Arc<dyn ChunkingAlgorithm>,
364 ) -> Self {
365 let buffer_size = config.ingestion_max_records.max(config.batch_size).max(2);
366 let using_config_triplet_recipes = !config.recipes.is_empty();
367 let using_config_text_recipes = !config.text_recipes.is_empty();
368 let triplet_recipes = if using_config_triplet_recipes {
369 config.recipes.clone()
370 } else {
371 Vec::new()
372 };
373 let text_recipes = if using_config_text_recipes {
374 config.text_recipes.clone()
375 } else if !triplet_recipes.is_empty() {
376 Self::build_derived_text_recipes(&triplet_recipes)
377 } else {
378 Vec::new()
379 };
380 let ingestion = IngestionManager::new(buffer_size, config.clone());
381 let epoch_backend = Some(Arc::clone(&split_store) as Arc<dyn EpochStateStore>);
382 let epoch_tracker = EpochTracker::new(
383 true,
384 epoch_backend,
385 derive_epoch_seed(config.seed, EPOCH_SEED_OFFSET),
386 );
387 let mut sampler = Self {
388 rng: DeterministicRng::new(config.seed),
389 config,
390 chunker,
391 split_store,
392 ingestion,
393 records: IndexMap::new(),
394 triplet_recipes,
395 text_recipes,
396 source_triplet_recipes: HashMap::new(),
397 sources_with_long_sections: HashSet::new(),
398 source_text_recipes: HashMap::new(),
399 using_config_triplet_recipes,
400 using_config_text_recipes,
401 last_observed_ingest: 0,
402 epoch_tracker,
403 chunk_cursors: HashMap::new(),
404 role_cursors: HashMap::new(),
405 negative_backend: {
406 #[cfg(feature = "bm25-mining")]
407 {
408 Box::new(backends::Bm25Backend::new())
409 }
410 #[cfg(not(feature = "bm25-mining"))]
411 {
412 Box::new(backends::DefaultBackend)
413 }
414 },
415 chunk_index: HashMap::new(),
416 source_order: Vec::new(),
417 source_cycle_idx: 0,
418 source_state_loaded: false,
419 ingestion_cursors_loaded: false,
420 source_state_dirty: false,
421 source_record_indices: HashMap::new(),
422 source_record_cursors: HashMap::new(),
423 emitted_text_hashes: HashMap::new(),
424 triplet_recipe_rr_idx: 0,
425 text_recipe_rr_idx: 0,
426 epoch: 0,
427 source_wrapped: HashMap::new(),
428 };
429 if !sampler.using_config_text_recipes {
430 sampler.rebuild_derived_text_recipes();
431 }
432 sampler
433 }
434
435 fn text_recipes(&self) -> &[TextRecipe] {
436 &self.text_recipes
437 }
438
439 fn epoch_seed(&self) -> u64 {
442 derive_epoch_seed(self.config.seed, self.epoch)
443 }
444
445 fn register_source(
446 &mut self,
447 source: Box<dyn DataSource + 'static>,
448 ) -> Result<(), SamplerError> {
449 let source_id = source.id().to_string();
450 if !self.using_config_triplet_recipes {
451 let triplets = source.default_triplet_recipes();
452 if !triplets.is_empty() {
453 self.source_triplet_recipes
454 .insert(source_id.clone(), triplets.clone());
455 if !self.using_config_text_recipes {
456 let derived = Self::build_derived_text_recipes(&triplets);
457 self.source_text_recipes
458 .insert(source_id.clone(), derived.clone());
459 self.extend_text_recipes_unique(&derived);
460 }
461 }
462 }
463 self.ingestion.register_source(source)?;
464 Ok(())
465 }
466
467 fn set_epoch(&mut self, epoch: u64) -> Result<(), SamplerError> {
468 self.epoch_tracker.ensure_loaded()?;
469 self.epoch_tracker.force_epoch(epoch);
470 self.epoch = epoch;
471 self.ingestion.set_epoch(epoch);
472 self.ingestion.reset_stream_cursors();
473 self.ingestion.reset_epoch_step();
476 self.source_record_cursors.clear();
477 self.source_cycle_idx = 0;
478 for source in &self.source_order {
479 self.source_wrapped.insert(source.clone(), false);
480 }
481 self.rebuild_source_index()?;
482 self.source_state_dirty = self.source_order.len() > 1;
483 Ok(())
484 }
485
486 fn next_chunk_from_pool(
487 &mut self,
488 record_id: &str,
489 section_idx: usize,
490 pool: Vec<RecordChunk>,
491 ) -> Option<RecordChunk> {
492 if pool.is_empty() {
493 return None;
494 }
495 let key = (record_id.to_string(), section_idx);
496 if !self.chunk_cursors.contains_key(&key) {
497 let cursor_key = format!("{}::{}", record_id, section_idx);
503 let start = (stable_hash_str(self.epoch_seed(), &cursor_key) as usize) % pool.len();
504 self.chunk_cursors.insert(key.clone(), start);
505 }
506 let cursor = self.chunk_cursors.entry(key).or_insert(0);
507 if *cursor >= pool.len() {
508 *cursor = 0;
509 }
510 let chunk = pool.get(*cursor).cloned();
511 *cursor = (*cursor + 1) % pool.len();
512 chunk
513 }
514
515 fn prune_cursor_state(&mut self) {
516 if self.chunk_cursors.is_empty()
517 && self.role_cursors.is_empty()
518 && self.negative_backend.cursors_empty()
519 {
520 return;
521 }
522 let valid_ids: HashSet<RecordId> = self.records.keys().cloned().collect();
523 self.chunk_cursors
524 .retain(|(record_id, _), _| valid_ids.contains(record_id));
525 self.role_cursors
526 .retain(|(record_id, _), _| valid_ids.contains(record_id));
527 self.negative_backend.prune_cursors(&valid_ids);
528 }
529
530 fn rebuild_chunk_index(&mut self) {
531 self.chunk_index.clear();
532 for record in self.records.values() {
533 self.chunk_index
534 .insert(record.id.clone(), record.id.clone());
535 }
536 }
537
538 fn rebuild_source_index(&mut self) -> Result<(), SamplerError> {
539 self.source_record_indices.clear();
540 let mut label_cache: HashMap<RecordId, SplitLabel> = HashMap::new();
541 let allowed = self.allowed_target_splits();
542 let allowed_set: HashSet<SplitLabel> = allowed.into_iter().collect();
543 for (idx, record) in self.records.values().enumerate() {
544 let label = if let Some(label) = label_cache.get(&record.id) {
545 *label
546 } else {
547 let label = match self.split_store.label_for(&record.id) {
548 Some(label) => label,
549 None => self.split_store.ensure(record.id.clone())?,
550 };
551 label_cache.insert(record.id.clone(), label);
552 label
553 };
554 if !allowed_set.contains(&label) {
555 continue;
556 }
557 self.source_record_indices
558 .entry(record.source.clone())
559 .or_default()
560 .push(idx);
561 }
562
563 let shuffle_seed = self.epoch_seed();
564 for indices in self.source_record_indices.values_mut() {
565 indices.sort_by_key(|idx| {
566 self.records
567 .get_index(*idx)
568 .map(|(_, record)| stable_hash_str(shuffle_seed, &record.id))
569 .unwrap_or(0)
570 });
571 }
572
573 self.source_order = self.source_record_indices.keys().cloned().collect();
574 self.source_order.sort();
575 self.refresh_source_wrapped();
576
577 self.source_record_cursors
578 .retain(|source, _| self.source_record_indices.contains_key(source));
579 if self.source_state_loaded {
580 if self.source_order.is_empty() {
581 self.source_cycle_idx = 0;
582 }
583 self.source_state_dirty = self.source_order.len() > 1;
584 }
585 Ok(())
586 }
587
588 fn refresh_source_wrapped(&mut self) {
589 self.source_wrapped.clear();
590 for source in &self.source_order {
591 let len = self
592 .source_record_indices
593 .get(source)
594 .map(|items| items.len())
595 .unwrap_or(0);
596 if len == 0 {
597 self.source_wrapped.insert(source.clone(), false);
598 continue;
599 }
600 let cursor = self.source_record_cursors.get(source).copied().unwrap_or(0);
601 let wrapped = cursor > 0 && cursor % len == 0;
602 self.source_wrapped.insert(source.clone(), wrapped);
603 }
604 }
605
606 fn shuffled_source_cycle(&self, cycle: u64) -> Vec<SourceId> {
607 let mut sources = self.source_order.clone();
608 let seed = self.epoch_seed() ^ cycle;
609 sources.sort_by_key(|source| stable_hash_str(seed, source));
610 sources
611 }
612
613 fn ensure_source_state(&mut self) -> Result<(), SamplerError> {
614 if self.source_state_loaded {
615 return Ok(());
616 }
617 let persisted = self.split_store.load_sampler_state()?;
618 self.source_cycle_idx = persisted
619 .as_ref()
620 .map(|state| state.source_cycle_idx as usize)
621 .unwrap_or(0);
622 if let Some(state) = persisted {
623 for (source, cursor) in state.source_record_cursors {
624 if self.source_record_indices.contains_key(&source) {
625 self.source_record_cursors.insert(source, cursor as usize);
626 }
627 }
628 self.epoch = state.epoch;
629 self.ingestion.set_epoch(state.epoch);
630 self.rng = DeterministicRng::from_state(state.rng_state);
631 self.triplet_recipe_rr_idx = state.triplet_recipe_rr_idx as usize;
632 self.text_recipe_rr_idx = state.text_recipe_rr_idx as usize;
633 }
634 self.refresh_source_wrapped();
635 self.source_state_loaded = true;
636 self.source_state_dirty = true;
637 Ok(())
638 }
639
640 fn persist_source_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
641 if !self.source_state_loaded {
642 return Ok(());
643 }
644 let state = PersistedSamplerState {
645 source_cycle_idx: self.source_cycle_idx as u64,
646 source_record_cursors: self
647 .source_record_cursors
648 .iter()
649 .map(|(source, cursor)| (source.clone(), *cursor as u64))
650 .collect(),
651 epoch: self.epoch,
652 epoch_step: self.ingestion.epoch_step(),
653 rng_state: self.rng.state(),
654 triplet_recipe_rr_idx: self.triplet_recipe_rr_idx as u64,
655 text_recipe_rr_idx: self.text_recipe_rr_idx as u64,
656 source_stream_cursors: self.ingestion.snapshot_cursors(),
657 };
658 self.split_store.save_sampler_state(&state, save_to)?;
659 self.source_state_dirty = false;
660 Ok(())
661 }
662
663 fn rebuild_derived_text_recipes(&mut self) {
664 if self.using_config_text_recipes {
665 return;
666 }
667 if self.triplet_recipes.is_empty() {
668 self.text_recipes.clear();
669 } else {
670 self.text_recipes = Self::build_derived_text_recipes(&self.triplet_recipes);
671 }
672 }
673
674 fn extend_text_recipes_unique(&mut self, recipes: &[TextRecipe]) {
675 for recipe in recipes {
676 if self
677 .text_recipes
678 .iter()
679 .any(|existing| existing.name == recipe.name)
680 {
681 continue;
682 }
683 self.text_recipes.push(recipe.clone());
684 }
685 }
686
687 fn configured_triplet_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TripletRecipe] {
688 if self.using_config_triplet_recipes {
689 return &self.triplet_recipes;
690 }
691 self.source_triplet_recipes
692 .get(source)
693 .map(|recipes| recipes.as_slice())
694 .unwrap_or(&[])
695 }
696
697 fn contains_auto_chunk_pair_recipe(recipes: &[TripletRecipe]) -> bool {
699 recipes
700 .iter()
701 .any(|recipe| recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME)
702 }
703
704 fn source_supports_chunk_pair_recipe(&self, source: &str) -> bool {
705 if self.config.chunking.max_window_tokens == 0 {
706 return false;
707 }
708 self.sources_with_long_sections.contains(source)
709 }
710
711 fn should_auto_inject_chunk_pair_recipe(
717 &self,
718 source: &str,
719 recipes: &[TripletRecipe],
720 ) -> bool {
721 self.source_supports_chunk_pair_recipe(source)
722 && !Self::contains_auto_chunk_pair_recipe(recipes)
723 }
724
725 fn source_chunk_pair_recipe() -> TripletRecipe {
736 TripletRecipe {
737 name: Cow::Borrowed(AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME),
738 anchor: Selector::Role(SectionRole::Context),
739 positive_selector: Selector::Role(SectionRole::Context),
740 negative_selector: Selector::Role(SectionRole::Context),
741 negative_strategy: NegativeStrategy::WrongArticle,
742 weight: 1.0,
743 instruction: None,
744 allow_same_anchor_positive: false,
745 }
746 }
747
748 fn resolve_source_triplet_plan(&self, source: &str) -> (Vec<TripletRecipe>, bool) {
760 let mut recipes = self.configured_triplet_recipes_for_source(source).to_vec();
761 let mut auto_injected = false;
762 if self.should_auto_inject_chunk_pair_recipe(source, &recipes) {
763 recipes.push(Self::source_chunk_pair_recipe());
764 auto_injected = true;
765 }
766 (recipes, auto_injected)
767 }
768
769 #[cfg(test)]
770 fn triplet_recipes_for_source(&self, source: &str) -> Vec<TripletRecipe> {
771 self.resolve_source_triplet_plan(source).0
772 }
773
774 fn triplet_recipe_count_for_source(&self, source: &str) -> usize {
775 let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
776 recipes.len()
777 }
778
779 fn text_recipes_for_source<'a>(&'a self, source: &str) -> &'a [TextRecipe] {
780 if self.using_config_text_recipes || self.using_config_triplet_recipes {
781 return &self.text_recipes;
782 }
783 self.source_text_recipes
784 .get(source)
785 .map(|recipes| recipes.as_slice())
786 .unwrap_or(&[])
787 }
788
789 fn recipe_order_weighted_shuffled(
800 &mut self,
801 weights: &[f32],
802 rng: &mut DeterministicRng,
803 ) -> Vec<usize> {
804 weighted_recipe_order(weights, rng)
805 }
806
807 fn recipe_order_weighted_cycled(
813 &mut self,
814 weights: &[f32],
815 rr_idx: usize,
816 rng: &mut DeterministicRng,
817 ) -> Vec<usize> {
818 let base = self.recipe_order_weighted_shuffled(weights, rng);
819 if base.is_empty() {
820 return base;
821 }
822 let start = rr_idx % base.len();
823 let mut order = Vec::with_capacity(base.len());
824 order.extend_from_slice(&base[start..]);
825 order.extend_from_slice(&base[..start]);
826 order
827 }
828
829 fn text_recipe_order_weighted_shuffled(
833 &mut self,
834 weights: &[f32],
835 rng: &mut DeterministicRng,
836 ) -> Vec<usize> {
837 weighted_recipe_order(weights, rng)
838 }
839
840 fn text_recipe_order_weighted_cycled(
841 &mut self,
842 weights: &[f32],
843 rr_idx: usize,
844 rng: &mut DeterministicRng,
845 ) -> Vec<usize> {
846 let base = self.text_recipe_order_weighted_shuffled(weights, rng);
847 if base.is_empty() {
848 return base;
849 }
850 let start = rr_idx % base.len();
851 let mut order = Vec::with_capacity(base.len());
852 order.extend_from_slice(&base[start..]);
853 order.extend_from_slice(&base[..start]);
854 order
855 }
856
857 fn allowed_target_splits(&self) -> Vec<SplitLabel> {
858 self.config.allowed_splits.clone()
859 }
860
861 fn ensure_split_allowed(&self, split: SplitLabel) -> Result<(), SamplerError> {
862 let allowed = self.allowed_target_splits();
863 if allowed.contains(&split) {
864 return Ok(());
865 }
866 Err(SamplerError::Configuration(format!(
867 "requested split {:?} is not in allowed_splits {:?}",
868 split, allowed
869 )))
870 }
871
872 fn ensure_split_has_records(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
873 let records_by_split = self.records_by_split()?;
874 if records_by_split
875 .get(&target_split)
876 .map(|records| !records.is_empty())
877 .unwrap_or(false)
878 {
879 return Ok(());
880 }
881 Err(SamplerError::Exhausted(
882 "no records available for target split".into(),
883 ))
884 }
885
886 fn records_by_split(
887 &self,
888 ) -> Result<HashMap<SplitLabel, Vec<(RecordId, SourceId)>>, SamplerError> {
889 let mut map: HashMap<SplitLabel, Vec<(RecordId, SourceId)>> = HashMap::new();
890 let mut label_cache: HashMap<RecordId, SplitLabel> = HashMap::new();
891 for (chunk_id, record_id) in &self.chunk_index {
892 let Some(record) = self.records.get(record_id) else {
893 continue;
894 };
895 let label = if let Some(label) = label_cache.get(record_id) {
896 *label
897 } else {
898 let label = match self.split_store.label_for(record_id) {
899 Some(label) => label,
900 None => self.split_store.ensure(record_id.clone())?,
901 };
902 label_cache.insert(record_id.clone(), label);
903 label
904 };
905 map.entry(label)
906 .or_default()
907 .push((chunk_id.clone(), record.source.clone()));
908 }
909 Ok(map)
910 }
911
912 fn choose_anchor_record(
913 &mut self,
914 source: Option<&str>,
915 split: SplitLabel,
916 ) -> Option<Arc<DataRecord>> {
917 if let Some(source) = source {
918 let indices = self.source_record_indices.get(source)?;
919 if indices.is_empty() {
920 return None;
921 }
922 let mut cursor = *self.source_record_cursors.get(source).unwrap_or(&0);
923 let cycle = cursor / indices.len();
924 let offset_seed = self.epoch_seed() ^ (cycle as u64);
925 let offset = (stable_hash_str(offset_seed, source) as usize) % indices.len();
926 let mut wrapped = false;
927 let mut selected: Option<Arc<DataRecord>> = None;
928 for _ in 0..indices.len() {
929 let pos = (cursor % indices.len()).saturating_add(offset) % indices.len();
930 let idx = indices[pos];
931 cursor = cursor.saturating_add(1);
932 if cursor.is_multiple_of(indices.len()) {
933 wrapped = true;
934 }
935 if let Some((_, record)) = self.records.get_index(idx) {
936 if self.split_store.label_for(&record.id) != Some(split) {
937 continue;
938 }
939 selected = Some(Arc::clone(record));
940 break;
941 }
942 }
943 self.source_record_cursors
944 .insert(source.to_string(), cursor);
945 if wrapped {
946 self.mark_source_wrapped(source);
947 }
948 return selected;
949 }
950 while let Some(chunk_id) = self.epoch_tracker.next_record(split) {
951 if let Some(record_id) = self.chunk_index.get(&chunk_id)
952 && let Some(record) = self.records.get(record_id)
953 {
954 return Some(Arc::clone(record));
955 }
956 }
957 None
958 }
959
960 fn save_sampler_state(&mut self, save_to: Option<&Path>) -> Result<(), SamplerError> {
961 if self.epoch_tracker.is_enabled() {
962 self.epoch_tracker.persist()?;
963 }
964 self.persist_source_state(save_to)?;
965 Ok(())
966 }
967
968 fn mark_source_wrapped(&mut self, source: &str) {
969 self.source_wrapped.insert(source.to_string(), true);
970 if self.source_order.is_empty() {
971 return;
972 }
973 let all_wrapped = self
974 .source_order
975 .iter()
976 .all(|name| self.source_wrapped.get(name).copied().unwrap_or(false));
977 if all_wrapped {
978 self.advance_epoch();
979 }
980 }
981
982 fn advance_epoch(&mut self) {
983 self.epoch = self.epoch.saturating_add(1);
990 self.ingestion.set_epoch(self.epoch);
991 self.source_record_cursors.clear();
998 self.source_cycle_idx = 0;
999 for source in &self.source_order {
1002 self.source_wrapped.insert(source.clone(), false);
1003 }
1004 let _ = self.rebuild_source_index();
1005 self.source_state_dirty = self.source_order.len() > 1;
1006 }
1007
1008 fn select_temporal_neighbor(
1009 &'_ self,
1010 record: &DataRecord,
1011 offset_days: i32,
1012 ) -> Option<Arc<DataRecord>> {
1013 let target = record.created_at + Duration::days(offset_days.into());
1014 let key = record.taxonomy.first().cloned();
1015 let record_split = self.split_store.label_for(&record.id)?;
1016 self.records
1017 .values()
1018 .filter(|candidate| {
1019 candidate.id != record.id
1020 && self
1021 .split_store
1022 .label_for(&candidate.id)
1023 .map(|label| label == record_split)
1024 .unwrap_or(false)
1025 && (candidate.source == record.source
1026 || key
1027 .as_ref()
1028 .zip(candidate.taxonomy.first())
1029 .map(|(a, b)| a == b)
1030 .unwrap_or(false))
1031 })
1032 .min_by_key(|candidate| (candidate.created_at - target).num_seconds().abs())
1033 .cloned()
1034 }
1035
1036 fn select_negative_record(
1037 &self,
1038 anchor_record: &DataRecord,
1039 strategy: &NegativeStrategy,
1040 anchor_query_text: Option<&str>,
1041 rng: &mut dyn rand::RngCore,
1042 ) -> Option<(Arc<DataRecord>, bool)> {
1043 let anchor_split = self.split_store.label_for(&anchor_record.id)?;
1044
1045 let in_anchor_split = |candidate: &DataRecord| {
1046 self.split_store
1047 .label_for(&candidate.id)
1048 .map(|label| label == anchor_split)
1049 .unwrap_or(false)
1050 };
1051
1052 match strategy {
1053 NegativeStrategy::WrongArticle => {
1054 let anchor_date =
1055 taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1056 let mut same_date: Vec<Arc<DataRecord>> = self
1057 .records
1058 .values()
1059 .filter(|candidate| {
1060 candidate.source == anchor_record.source
1061 && candidate.id != anchor_record.id
1062 && in_anchor_split(candidate)
1063 })
1064 .filter(|candidate| {
1065 anchor_date
1066 .as_deref()
1067 .zip(taxonomy_value(candidate, META_FIELD_DATE))
1068 .map(|(a, b)| a == b)
1069 .unwrap_or(false)
1070 })
1071 .cloned()
1072 .collect();
1073 if same_date.is_empty() {
1074 same_date = self
1075 .records
1076 .values()
1077 .filter(|candidate| {
1078 candidate.source == anchor_record.source
1079 && candidate.id != anchor_record.id
1080 && in_anchor_split(candidate)
1081 })
1082 .cloned()
1083 .collect();
1084 }
1085 if !same_date.is_empty() {
1086 return self.negative_backend.choose_negative(
1087 anchor_record,
1088 anchor_split,
1089 same_date,
1090 false,
1091 anchor_query_text,
1092 rng,
1093 );
1094 }
1095 let pool = self
1096 .records
1097 .values()
1098 .filter(|candidate| {
1099 candidate.id != anchor_record.id && in_anchor_split(candidate)
1100 })
1101 .cloned()
1102 .collect::<Vec<_>>();
1103 self.negative_backend.choose_negative(
1104 anchor_record,
1105 anchor_split,
1106 pool,
1107 true,
1108 anchor_query_text,
1109 rng,
1110 )
1111 }
1112 NegativeStrategy::WrongPublicationDate => {
1113 let anchor_date =
1114 taxonomy_value(anchor_record, META_FIELD_DATE).map(|d| d.to_string());
1115 let pool: Vec<Arc<DataRecord>> = self
1116 .records
1117 .values()
1118 .filter(|candidate| {
1119 candidate.source == anchor_record.source
1120 && candidate.id != anchor_record.id
1121 && in_anchor_split(candidate)
1122 })
1123 .filter(|candidate| {
1124 match (
1125 anchor_date.as_deref(),
1126 taxonomy_value(candidate, META_FIELD_DATE),
1127 ) {
1128 (Some(anchor), Some(candidate_date)) => anchor != candidate_date,
1129 (Some(_), None) => true,
1130 (None, Some(_)) => true,
1131 (None, None) => false,
1132 }
1133 })
1134 .cloned()
1135 .collect();
1136 if pool.is_empty() {
1137 let fallback_pool = self
1140 .records
1141 .values()
1142 .filter(|candidate| {
1143 candidate.id != anchor_record.id && in_anchor_split(candidate)
1144 })
1145 .cloned()
1146 .collect::<Vec<_>>();
1147
1148 return self.negative_backend.choose_negative(
1149 anchor_record,
1150 anchor_split,
1151 fallback_pool,
1152 true,
1153 anchor_query_text,
1154 rng,
1155 );
1156 }
1157
1158 self.negative_backend.choose_negative(
1159 anchor_record,
1160 anchor_split,
1161 pool,
1162 false,
1163 anchor_query_text,
1164 rng,
1165 )
1166 }
1167 NegativeStrategy::QuestionAnswerMismatch => {
1168 let pool: Vec<Arc<DataRecord>> = self
1169 .records
1170 .values()
1171 .filter(|candidate| {
1172 candidate.source == anchor_record.source
1173 && candidate.id != anchor_record.id
1174 && in_anchor_split(candidate)
1175 })
1176 .cloned()
1177 .collect();
1178 if pool.is_empty() {
1179 let fallback_pool = self
1182 .records
1183 .values()
1184 .filter(|candidate| {
1185 candidate.id != anchor_record.id && in_anchor_split(candidate)
1186 })
1187 .cloned()
1188 .collect::<Vec<_>>();
1189
1190 return self.negative_backend.choose_negative(
1191 anchor_record,
1192 anchor_split,
1193 fallback_pool,
1194 true,
1195 anchor_query_text,
1196 rng,
1197 );
1198 }
1199
1200 self.negative_backend.choose_negative(
1201 anchor_record,
1202 anchor_split,
1203 pool,
1204 false,
1205 anchor_query_text,
1206 rng,
1207 )
1208 }
1209 NegativeStrategy::SameRecord => {
1210 Some((Arc::new(anchor_record.clone()), false))
1216 }
1217 }
1218 }
1219
1220 fn is_auto_chunk_pair_recipe(recipe: &TripletRecipe) -> bool {
1222 recipe.name.as_ref() == AUTO_INJECTED_LONG_SECTION_CHUNK_PAIR_RECIPE_NAME
1223 }
1224
1225 fn select_anchor_positive_pair(
1229 &mut self,
1230 record: &DataRecord,
1231 anchor_selector: &Selector,
1232 positive_selector: &Selector,
1233 enforce_window_pair: bool,
1234 ) -> Option<(RecordChunk, RecordChunk)> {
1235 let mut anchor_chunk = self.select_chunk(record, anchor_selector)?;
1236 let mut positive_chunk = self.select_chunk(record, positive_selector)?;
1237 if anchor_selector == positive_selector {
1238 let mut retries = 0usize;
1239 while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1240 && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1241 {
1242 let Some(redraw_anchor) = self.select_chunk(record, anchor_selector) else {
1243 break;
1244 };
1245 let Some(redraw_positive) = self.select_chunk(record, positive_selector) else {
1246 break;
1247 };
1248 anchor_chunk = redraw_anchor;
1249 positive_chunk = redraw_positive;
1250 retries += 1;
1251 }
1252 if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1253 return None;
1254 }
1255 }
1256 Some((anchor_chunk, positive_chunk))
1257 }
1258
1259 fn select_distinct_window_pair_for_auto_recipe(
1269 &mut self,
1270 recipe: &TripletRecipe,
1271 record: &DataRecord,
1272 ) -> Option<(RecordChunk, RecordChunk)> {
1273 if recipe.anchor != recipe.positive_selector {
1274 return None;
1275 }
1276 self.select_anchor_positive_pair(record, &recipe.anchor, &recipe.positive_selector, true)
1277 }
1278
1279 fn record_has_at_least_two_window_chunks_for_selector(
1284 &self,
1285 record: &DataRecord,
1286 selector: &Selector,
1287 ) -> bool {
1288 let section_indices: Vec<usize> = match selector {
1289 Selector::Role(role) => record
1290 .sections
1291 .iter()
1292 .enumerate()
1293 .filter(|(_, section)| roles_match(role, §ion.role))
1294 .map(|(idx, _)| idx)
1295 .collect(),
1296 Selector::Paragraph(idx) => {
1297 if *idx < record.sections.len() {
1298 vec![*idx]
1299 } else {
1300 Vec::new()
1301 }
1302 }
1303 Selector::Random => (0..record.sections.len()).collect(),
1304 Selector::TemporalOffset(_) => return false,
1305 };
1306
1307 let mut window_count = 0usize;
1308 for section_idx in section_indices {
1309 let Some(section) = record.sections.get(section_idx) else {
1310 continue;
1311 };
1312 let chunks = self.materialize_chunks(record, section_idx, section);
1313 window_count += chunks
1314 .iter()
1315 .filter(|chunk| matches!(chunk.view, ChunkView::Window { .. }))
1316 .count();
1317 if window_count >= 2 {
1318 return true;
1319 }
1320 }
1321 false
1322 }
1323
1324 fn build_triplet_with_selector_pair_policy(
1326 &mut self,
1327 recipe: &TripletRecipe,
1328 record: &DataRecord,
1329 enforce_window_pair: bool,
1330 rng: &mut DeterministicRng,
1331 ) -> Option<SampleTriplet> {
1332 let (mut anchor_chunk, mut positive_chunk) = self.select_anchor_positive_pair(
1333 record,
1334 &recipe.anchor,
1335 &recipe.positive_selector,
1336 enforce_window_pair,
1337 )?;
1338 let anchor_raw_text = anchor_chunk.text.clone();
1339 self.decorate_chunk(record, &mut anchor_chunk, rng);
1340 self.decorate_chunk(record, &mut positive_chunk, rng);
1341 self.finalize_triplet_with_negative(
1344 recipe,
1345 record,
1346 anchor_chunk,
1347 positive_chunk,
1348 &anchor_raw_text,
1349 rng,
1350 )
1351 }
1352
1353 fn make_auto_chunk_pair_triplet_with_anchor(
1363 &mut self,
1364 recipe: &TripletRecipe,
1365 record: &DataRecord,
1366 rng: &mut DeterministicRng,
1367 ) -> Option<SampleTriplet> {
1368 if !self.record_has_at_least_two_window_chunks_for_selector(record, &recipe.anchor) {
1369 return None;
1370 }
1371 let (mut anchor_chunk, mut positive_chunk) =
1372 self.select_distinct_window_pair_for_auto_recipe(recipe, record)?;
1373 let anchor_raw_text = anchor_chunk.text.clone();
1374 self.decorate_chunk(record, &mut anchor_chunk, rng);
1375 self.decorate_chunk(record, &mut positive_chunk, rng);
1376 self.finalize_triplet_with_negative(
1377 recipe,
1378 record,
1379 anchor_chunk,
1380 positive_chunk,
1381 &anchor_raw_text,
1382 rng,
1383 )
1384 }
1385
1386 fn make_standard_triplet_with_anchor(
1387 &mut self,
1388 recipe: &TripletRecipe,
1389 record: &DataRecord,
1390 rng: &mut DeterministicRng,
1391 ) -> Option<SampleTriplet> {
1392 self.build_triplet_with_selector_pair_policy(recipe, record, false, rng)
1393 }
1394
1395 fn finalize_triplet_with_negative(
1414 &mut self,
1415 recipe: &TripletRecipe,
1416 record: &DataRecord,
1417 anchor_chunk: RecordChunk,
1418 positive_chunk: RecordChunk,
1419 anchor_raw_text: &str,
1420 rng: &mut DeterministicRng,
1421 ) -> Option<SampleTriplet> {
1422 let (negative_record, fallback_used) = self.select_negative_record(
1423 record,
1424 &recipe.negative_strategy,
1425 Some(anchor_raw_text),
1426 rng,
1427 )?;
1428 let mut negative_chunk = self.select_chunk(&negative_record, &recipe.negative_selector)?;
1429 self.decorate_chunk(&negative_record, &mut negative_chunk, rng);
1430
1431 let (anchor_chunk, positive_chunk) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0 {
1433 (positive_chunk, anchor_chunk)
1434 } else {
1435 (anchor_chunk, positive_chunk)
1436 };
1437
1438 if (!recipe.allow_same_anchor_positive && anchor_chunk.text == positive_chunk.text)
1449 || negative_chunk.text == positive_chunk.text
1450 || negative_chunk.text == anchor_chunk.text
1451 {
1452 return None;
1453 }
1454
1455 let chunk_weight =
1456 self.triplet_chunk_weight(recipe, &anchor_chunk, &positive_chunk, &negative_chunk);
1457 let weight = recipe.weight * chunk_weight;
1458 let recipe_name = if fallback_used {
1459 format!("{}_fallback_same_split", recipe.name)
1460 } else {
1461 recipe.name.to_string()
1462 };
1463 Some(SampleTriplet {
1464 recipe: recipe_name,
1465 anchor: anchor_chunk,
1466 positive: positive_chunk,
1467 negative: negative_chunk,
1468 weight,
1469 instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1470 })
1471 }
1472
1473 fn make_triplet_with_anchor(
1474 &mut self,
1475 recipe: &TripletRecipe,
1476 record: &DataRecord,
1477 rng: &mut DeterministicRng,
1478 ) -> Option<SampleTriplet> {
1479 if Self::is_auto_chunk_pair_recipe(recipe) {
1480 return self.make_auto_chunk_pair_triplet_with_anchor(recipe, record, rng);
1481 }
1482 self.make_standard_triplet_with_anchor(recipe, record, rng)
1483 }
1484
1485 fn make_text_sample_for_split(
1486 &mut self,
1487 recipe: &TextRecipe,
1488 source: Option<&str>,
1489 split: SplitLabel,
1490 rng: &mut DeterministicRng,
1491 ) -> Option<TextSample> {
1492 let record = self.choose_anchor_record(source, split)?;
1493 let mut chunk = self.select_chunk(&record, &recipe.selector)?;
1494 self.decorate_chunk(&record, &mut chunk, rng);
1495 let weight = recipe.weight * self.chunk_weight(&chunk);
1496 Some(TextSample {
1497 recipe: recipe.name.to_string(),
1498 chunk,
1499 weight,
1500 instruction: recipe.instruction.as_ref().map(|s| s.to_string()),
1501 })
1502 }
1503
1504 fn chunk_weight(&self, chunk: &RecordChunk) -> f32 {
1505 chunk_weight(&self.config.chunking, chunk)
1506 }
1507
1508 fn triplet_chunk_weight(
1509 &self,
1510 recipe: &TripletRecipe,
1511 anchor: &RecordChunk,
1512 positive: &RecordChunk,
1513 negative: &RecordChunk,
1514 ) -> f32 {
1515 let floor = self.config.chunking.chunk_weight_floor;
1516 let negative_weight = negative.quality.trust.clamp(0.0, 1.0).max(floor);
1517 if Self::is_auto_chunk_pair_recipe(recipe) {
1518 let pair_trust = ((anchor.quality.trust.clamp(0.0, 1.0)
1521 + positive.quality.trust.clamp(0.0, 1.0))
1522 / 2.0)
1523 .clamp(0.0, 1.0);
1524 let pair_weight = (chunk_proximity_score(anchor, positive) * pair_trust).max(floor);
1525 return (pair_weight + pair_weight + negative_weight) / 3.0;
1527 }
1528 let pair_proximity = chunk_proximity_score(anchor, positive);
1531 let anchor_weight = (self.chunk_weight(anchor) * pair_proximity).max(floor);
1532 let positive_weight = (self.chunk_weight(positive) * pair_proximity).max(floor);
1533 (anchor_weight + positive_weight + negative_weight) / 3.0
1534 }
1535
1536 fn decorate_chunk(
1537 &mut self,
1538 record: &DataRecord,
1539 chunk: &mut RecordChunk,
1540 rng: &mut DeterministicRng,
1541 ) {
1542 chunk.kvp_meta = record
1543 .meta_prefix
1544 .as_ref()
1545 .map(|s| s.all_metadata())
1546 .unwrap_or_default();
1547 if let Some(spec) = record.meta_prefix.as_ref()
1548 && let Some(prefix) = spec.sample(rng)
1549 {
1550 let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1551 let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1552 let total_tokens = prefix_tokens.len() + body_tokens.len();
1553 let max_window = self.config.chunking.max_window_tokens;
1554 if max_window > 0 && total_tokens > max_window {
1555 if prefix_tokens.len() >= max_window {
1556 chunk.text = prefix_tokens
1557 .into_iter()
1558 .take(max_window)
1559 .collect::<Vec<_>>()
1560 .join(" ");
1561 chunk.tokens_estimate = max_window;
1562 } else {
1563 let remaining = max_window - prefix_tokens.len();
1564 let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1565 chunk.text =
1566 format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1567 chunk.tokens_estimate = max_window;
1568 }
1569 } else {
1570 chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1571 chunk.tokens_estimate = total_tokens;
1572 }
1573 }
1574 }
1575
1576 fn select_chunk_parallel(
1580 &self,
1581 record: &DataRecord,
1582 selector: &Selector,
1583 rng: &mut DeterministicRng,
1584 ) -> Option<RecordChunk> {
1585 match selector {
1586 Selector::Role(role) => self.select_role_parallel(record, role, rng),
1587 Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1588 let pool = self.materialize_chunks(record, *idx, section);
1589 if pool.is_empty() {
1590 return None;
1591 }
1592 let i = rng.random_range(0..pool.len());
1593 pool.into_iter().nth(i)
1594 }),
1595 Selector::TemporalOffset(offset) => self
1596 .select_temporal_neighbor(record, *offset)
1597 .and_then(|neighbor| {
1598 self.select_role_parallel(&neighbor, &SectionRole::Context, rng)
1599 }),
1600 Selector::Random => {
1601 if record.sections.is_empty() {
1602 return None;
1603 }
1604 let idx = rng.random_range(0..record.sections.len());
1605 record.sections.get(idx).and_then(|section| {
1606 let pool = self.materialize_chunks(record, idx, section);
1607 if pool.is_empty() {
1608 return None;
1609 }
1610 let i = rng.random_range(0..pool.len());
1611 pool.into_iter().nth(i)
1612 })
1613 }
1614 }
1615 }
1616
1617 fn select_role_parallel(
1619 &self,
1620 record: &DataRecord,
1621 role: &SectionRole,
1622 rng: &mut DeterministicRng,
1623 ) -> Option<RecordChunk> {
1624 let indices: Vec<usize> = record
1625 .sections
1626 .iter()
1627 .enumerate()
1628 .filter(|(_, s)| roles_match(role, &s.role))
1629 .map(|(i, _)| i)
1630 .collect();
1631 if indices.is_empty() {
1632 return None;
1633 }
1634 let start = rng.random_range(0..indices.len());
1635 for offset in 0..indices.len() {
1636 let section_idx = indices[(start + offset) % indices.len()];
1637 let section = &record.sections[section_idx];
1638 let pool = self.materialize_chunks(record, section_idx, section);
1639 if !pool.is_empty() {
1640 let i = rng.random_range(0..pool.len());
1641 return pool.into_iter().nth(i);
1642 }
1643 }
1644 None
1645 }
1646
1647 fn decorate_chunk_parallel(
1649 &self,
1650 record: &DataRecord,
1651 chunk: &mut RecordChunk,
1652 rng: &mut DeterministicRng,
1653 ) {
1654 chunk.kvp_meta = record
1655 .meta_prefix
1656 .as_ref()
1657 .map(|s| s.all_metadata())
1658 .unwrap_or_default();
1659 if let Some(spec) = record.meta_prefix.as_ref()
1660 && let Some(prefix) = spec.sample(rng)
1661 {
1662 let body_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&chunk.text);
1663 let prefix_tokens: Vec<&str> = WhitespaceTokenizer.tokenize(&prefix);
1664 let total_tokens = prefix_tokens.len() + body_tokens.len();
1665 let max_window = self.config.chunking.max_window_tokens;
1666 if max_window > 0 && total_tokens > max_window {
1667 if prefix_tokens.len() >= max_window {
1668 chunk.text = prefix_tokens
1669 .into_iter()
1670 .take(max_window)
1671 .collect::<Vec<_>>()
1672 .join(" ");
1673 chunk.tokens_estimate = max_window;
1674 } else {
1675 let remaining = max_window - prefix_tokens.len();
1676 let trimmed_body: Vec<&str> = body_tokens.into_iter().take(remaining).collect();
1677 chunk.text =
1678 format!("{}{}{}", prefix, platform_newline(), trimmed_body.join(" "));
1679 chunk.tokens_estimate = max_window;
1680 }
1681 } else {
1682 chunk.text = format!("{}{}{}", prefix, platform_newline(), chunk.text);
1683 chunk.tokens_estimate = total_tokens;
1684 }
1685 }
1686 }
1687
1688 fn select_anchor_positive_parallel(
1690 &self,
1691 record: &DataRecord,
1692 anchor_selector: &Selector,
1693 positive_selector: &Selector,
1694 enforce_window_pair: bool,
1695 rng: &mut DeterministicRng,
1696 ) -> Option<(RecordChunk, RecordChunk)> {
1697 let anchor_chunk = self.select_chunk_parallel(record, anchor_selector, rng)?;
1698 let mut positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1699 if anchor_selector == positive_selector {
1700 let mut retries = 0usize;
1701 while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair)
1702 && retries < SAME_SELECTOR_PAIR_RETRY_LIMIT
1703 {
1704 positive_chunk = self.select_chunk_parallel(record, positive_selector, rng)?;
1705 retries += 1;
1706 }
1707 if !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, enforce_window_pair) {
1708 return None;
1709 }
1710 }
1711 Some((anchor_chunk, positive_chunk))
1713 }
1714
1715 fn select_anchor_positive_for_recipe(
1718 &self,
1719 recipe: &TripletRecipe,
1720 anchor_record: &DataRecord,
1721 rng: &mut DeterministicRng,
1722 ) -> Option<(RecordChunk, RecordChunk, String)> {
1723 if Self::is_auto_chunk_pair_recipe(recipe) {
1724 if !self
1725 .record_has_at_least_two_window_chunks_for_selector(anchor_record, &recipe.anchor)
1726 {
1727 return None;
1728 }
1729 let mut anchor_chunk =
1730 self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1731 let mut positive_chunk =
1732 self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1733 let mut tries = 0usize;
1734 while !same_selector_pair_is_valid(&anchor_chunk, &positive_chunk, true) {
1735 tries += 1;
1736 if tries >= SAME_SELECTOR_PAIR_RETRY_LIMIT {
1737 return None;
1738 }
1739 anchor_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1740 positive_chunk = self.select_chunk_parallel(anchor_record, &recipe.anchor, rng)?;
1741 }
1742 let anchor_raw_text = anchor_chunk.text.clone();
1743 self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1744 self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1745 return Some((anchor_chunk, positive_chunk, anchor_raw_text));
1746 }
1747 let (mut anchor_chunk, mut positive_chunk) = self.select_anchor_positive_parallel(
1748 anchor_record,
1749 &recipe.anchor,
1750 &recipe.positive_selector,
1751 false,
1752 rng,
1753 )?;
1754 let anchor_raw_text = anchor_chunk.text.clone();
1755 self.decorate_chunk_parallel(anchor_record, &mut anchor_chunk, rng);
1756 self.decorate_chunk_parallel(anchor_record, &mut positive_chunk, rng);
1757 Some((anchor_chunk, positive_chunk, anchor_raw_text))
1758 }
1759
1760 fn select_chunk(&mut self, record: &DataRecord, selector: &Selector) -> Option<RecordChunk> {
1761 match selector {
1762 Selector::Role(role) => self.select_by_role(record, role),
1763 Selector::Paragraph(idx) => record.sections.get(*idx).and_then(|section| {
1764 let pool = self.materialize_chunks(record, *idx, section);
1765 self.next_chunk_from_pool(&record.id, *idx, pool)
1766 }),
1767 Selector::TemporalOffset(offset) => self
1768 .select_temporal_neighbor(record, *offset)
1769 .and_then(|neighbor| self.select_by_role(&neighbor, &SectionRole::Context)),
1770 Selector::Random => {
1771 if record.sections.is_empty() {
1772 return None;
1773 }
1774 let idx = self.rng.random_range(0..record.sections.len());
1775 record.sections.get(idx).and_then(|section| {
1776 let pool = self.materialize_chunks(record, idx, section);
1777 self.next_chunk_from_pool(&record.id, idx, pool)
1778 })
1779 }
1780 }
1781 }
1782
1783 fn select_by_role(&mut self, record: &DataRecord, role: &SectionRole) -> Option<RecordChunk> {
1784 let indices: Vec<usize> = record
1785 .sections
1786 .iter()
1787 .enumerate()
1788 .filter(|(_, section)| roles_match(role, §ion.role))
1789 .map(|(idx, _)| idx)
1790 .collect();
1791 if indices.is_empty() {
1792 return None;
1793 }
1794 let key = role_cursor_key(&record.id, role);
1795 let start_offset = self
1796 .role_cursors
1797 .get(&key)
1798 .and_then(|last_idx| indices.iter().position(|idx| idx == last_idx))
1799 .map(|pos| (pos + 1) % indices.len())
1800 .unwrap_or_else(|| {
1801 let seed_key = format!("{}::{}", key.0, key.1);
1805 (stable_hash_str(self.epoch_seed(), &seed_key) as usize) % indices.len()
1806 });
1807 for offset in 0..indices.len() {
1808 let section_idx = indices[(start_offset + offset) % indices.len()];
1809 let section = &record.sections[section_idx];
1810 let pool = self.materialize_chunks(record, section_idx, section);
1811 if let Some(chunk) = self.next_chunk_from_pool(&record.id, section_idx, pool) {
1812 self.role_cursors.insert(key.clone(), section_idx);
1813 return Some(chunk);
1814 }
1815 }
1816 None
1817 }
1818
1819 fn materialize_chunks(
1832 &self,
1833 record: &DataRecord,
1834 section_idx: usize,
1835 section: &RecordSection,
1836 ) -> Vec<RecordChunk> {
1837 self.chunker
1838 .materialize(&self.config.chunking, record, section_idx, section)
1839 }
1840
1841 fn build_derived_text_recipes(recipes: &[TripletRecipe]) -> Vec<TextRecipe> {
1842 let mut derived = Vec::new();
1843 for recipe in recipes {
1844 let base = recipe.name.as_ref();
1845 derived.push(TextRecipe {
1846 name: Cow::Owned(format!("{base}_anchor")),
1847 selector: recipe.anchor.clone(),
1848 weight: recipe.weight,
1849 instruction: None,
1850 });
1851 derived.push(TextRecipe {
1852 name: Cow::Owned(format!("{base}_positive")),
1853 selector: recipe.positive_selector.clone(),
1854 weight: recipe.weight,
1855 instruction: None,
1856 });
1857 derived.push(TextRecipe {
1858 name: Cow::Owned(format!("{base}_negative")),
1859 selector: recipe.negative_selector.clone(),
1860 weight: recipe.weight,
1861 instruction: None,
1862 });
1863 }
1864 derived
1865 }
1866
1867 fn record_has_long_anchor_or_context_section(&self, record: &DataRecord) -> bool {
1870 let window = self.config.chunking.max_window_tokens;
1871 if window == 0 {
1872 return false;
1873 }
1874 record.sections.iter().any(|section| {
1875 matches!(section.role, SectionRole::Anchor | SectionRole::Context)
1876 && WhitespaceTokenizer.token_count(§ion.text) > window
1877 })
1878 }
1879
1880 fn sync_records_from_cache(&mut self) -> Result<(), SamplerError> {
1881 let mut snapshot = self.ingestion.all_records_snapshot();
1882 snapshot.sort_by(|a, b| a.id.cmp(&b.id));
1883
1884 let old_ids: HashSet<String> = self.records.keys().cloned().collect();
1887 let incoming_ids: HashSet<String> = snapshot.iter().map(|r| r.id.clone()).collect();
1888 let evicted: HashSet<&String> = old_ids.difference(&incoming_ids).collect();
1889
1890 self.records.clear();
1895 self.sources_with_long_sections.clear();
1899 for evicted_id in &evicted {
1908 self.emitted_text_hashes.remove(*evicted_id);
1909 }
1910 self.negative_backend.on_sync_start();
1915 for record in snapshot {
1916 if self.split_store.label_for(&record.id).is_none() {
1917 self.split_store.ensure(record.id.clone())?;
1918 }
1919 if self.record_has_long_anchor_or_context_section(&record) {
1920 self.sources_with_long_sections
1922 .insert(record.source.clone());
1923 }
1924 self.records.insert(record.id.clone(), Arc::new(record));
1925 }
1926 self.prune_cursor_state();
1927 self.rebuild_chunk_index();
1928 self.rebuild_source_index()?;
1929 Ok(())
1930 }
1931
1932 fn ingest_internal_for_split(&mut self, target_split: SplitLabel) -> Result<(), SamplerError> {
1933 if !self.ingestion.has_sources() {
1934 return Ok(());
1935 }
1936 if !self.ingestion_cursors_loaded {
1937 if let Some(state) = self.split_store.load_sampler_state()? {
1938 self.ingestion.load_cursors(&state.source_stream_cursors);
1939 self.ingestion.set_epoch(state.epoch);
1940 self.ingestion.set_epoch_step(state.epoch_step);
1941 }
1942 self.ingestion_cursors_loaded = true;
1943 }
1944 if self.ingestion.all_caches_empty() {
1945 self.ingestion.refresh_all();
1946 } else {
1947 self.ingestion.advance(self.config.batch_size);
1948 }
1949 let observed = self.ingestion.total_ingest_count();
1950 if observed == self.last_observed_ingest && !self.records.is_empty() {
1951 return Ok(());
1952 }
1953 self.last_observed_ingest = observed;
1954 self.sync_records_from_cache()?;
1955 let max_window_tokens = self.config.chunking.max_window_tokens;
1956 self.negative_backend.on_records_refreshed(
1957 &self.records,
1958 max_window_tokens,
1959 &|id| self.split_store.label_for(id),
1960 self.ingestion.last_refreshed_sources(),
1961 );
1962 self.epoch_tracker.ensure_loaded()?;
1965 let records_by_split = self.records_by_split()?;
1966 self.epoch_tracker
1967 .reconcile(target_split, &records_by_split);
1968 self.ensure_source_state()?;
1969 Ok(())
1970 }
1971
1972 fn ensure_ingestion_cursors_loaded(&mut self) -> Result<(), SamplerError> {
1981 if !self.ingestion_cursors_loaded {
1982 if let Some(state) = self.split_store.load_sampler_state()? {
1983 self.ingestion.load_cursors(&state.source_stream_cursors);
1984 self.ingestion.set_epoch(state.epoch);
1985 self.ingestion.set_epoch_step(state.epoch_step);
1986 }
1987 self.ingestion_cursors_loaded = true;
1988 }
1989 Ok(())
1990 }
1991
1992 #[cfg(test)]
1993 fn ingest_internal(&mut self, split: SplitLabel) -> Result<(), SamplerError> {
1994 self.ingest_internal_for_split(split)
1995 }
1996
1997 fn ingest_internal_with_weights_for_split(
1998 &mut self,
1999 target_split: SplitLabel,
2000 weights: &HashMap<SourceId, f32>,
2001 ) -> Result<(), SamplerError> {
2002 if !self.ingestion.has_sources() {
2003 return Ok(());
2004 }
2005 if !self.ingestion_cursors_loaded {
2006 if let Some(state) = self.split_store.load_sampler_state()? {
2007 self.ingestion.load_cursors(&state.source_stream_cursors);
2008 self.ingestion.set_epoch(state.epoch);
2009 self.ingestion.set_epoch_step(state.epoch_step);
2010 }
2011 self.ingestion_cursors_loaded = true;
2012 }
2013 if self.ingestion.all_caches_empty() {
2014 self.ingestion.refresh_all_with_weights(weights)?;
2015 } else {
2016 self.ingestion
2017 .advance_with_weights(self.config.batch_size, weights)?;
2018 }
2019 let observed = self.ingestion.total_ingest_count();
2020 if observed == self.last_observed_ingest && !self.records.is_empty() {
2021 return Ok(());
2022 }
2023 self.last_observed_ingest = observed;
2024 self.sync_records_from_cache()?;
2025 let max_window_tokens = self.config.chunking.max_window_tokens;
2026 self.negative_backend.on_records_refreshed(
2027 &self.records,
2028 max_window_tokens,
2029 &|id| self.split_store.label_for(id),
2030 self.ingestion.last_refreshed_sources(),
2031 );
2032 self.epoch_tracker.ensure_loaded()?;
2033 let records_by_split = self.records_by_split()?;
2034 self.epoch_tracker
2035 .reconcile(target_split, &records_by_split);
2036 self.ensure_source_state()?;
2037 Ok(())
2038 }
2039
2040 fn ingest_with_weights_fallback(
2045 &mut self,
2046 target_split: SplitLabel,
2047 weights: Option<&HashMap<SourceId, f32>>,
2048 ) -> Result<(), SamplerError> {
2049 match weights {
2050 Some(weights)
2051 if !weights.is_empty()
2052 && !weights
2053 .values()
2054 .all(|&w| w == *weights.values().next().unwrap()) =>
2055 {
2056 self.ingest_internal_with_weights_for_split(target_split, weights)?
2057 }
2058 _ => self.ingest_internal_for_split(target_split)?,
2059 }
2060 Ok(())
2061 }
2062
2063 fn force_ingest_refresh_with_weights_for_split(
2064 &mut self,
2065 target_split: SplitLabel,
2066 weights: &HashMap<SourceId, f32>,
2067 ) -> Result<(), SamplerError> {
2068 if !self.ingestion.has_sources() {
2069 return Ok(());
2070 }
2071 if !self.ingestion_cursors_loaded {
2072 if let Some(state) = self.split_store.load_sampler_state()? {
2073 self.ingestion.load_cursors(&state.source_stream_cursors);
2074 self.ingestion.set_epoch(state.epoch);
2075 self.ingestion.set_epoch_step(state.epoch_step);
2076 }
2077 self.ingestion_cursors_loaded = true;
2078 }
2079 self.ingestion.force_refresh_all_with_weights(weights)?;
2080 self.last_observed_ingest = self.ingestion.total_ingest_count();
2081 self.sync_records_from_cache()?;
2082 let max_window_tokens = self.config.chunking.max_window_tokens;
2083 self.negative_backend.on_records_refreshed(
2084 &self.records,
2085 max_window_tokens,
2086 &|id| self.split_store.label_for(id),
2087 self.ingestion.last_refreshed_sources(),
2088 );
2089 self.epoch_tracker.ensure_loaded()?;
2090 let records_by_split = self.records_by_split()?;
2091 self.epoch_tracker
2092 .reconcile(target_split, &records_by_split);
2093 self.ensure_source_state()?;
2094 Ok(())
2095 }
2096
2097 fn sample_source_triplet_candidate(
2107 &mut self,
2108 source: &str,
2109 target_split: SplitLabel,
2110 recipe_orders: &mut HashMap<RecipeKey, Vec<usize>>,
2111 recipe_positions: &mut HashMap<RecipeKey, usize>,
2112 rng: &mut DeterministicRng,
2113 ) -> (Option<(TripletRecipe, SampleTriplet)>, usize) {
2114 let (recipes, _auto_injected) = self.resolve_source_triplet_plan(source);
2117 if recipes.is_empty() {
2118 return (None, 0);
2119 }
2120 if !recipe_orders.contains_key(source) {
2121 let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2122 let order =
2123 self.recipe_order_weighted_cycled(&recipe_weights, self.triplet_recipe_rr_idx, rng);
2124 recipe_orders.insert(source.to_string(), order);
2125 }
2126 let order = recipe_orders
2127 .get(source)
2128 .expect("recipe order missing for source");
2129 let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2130 let Some(anchor) = self.choose_anchor_record(Some(source), target_split) else {
2131 return (None, 0);
2132 };
2133
2134 let mut attempts = 0usize;
2135 for offset in 0..order.len() {
2136 let idx = order[(*pos + offset) % order.len()];
2137 attempts = attempts.saturating_add(1);
2138 let recipe = recipes[idx].clone();
2139 if let Some(sample) = self.make_triplet_with_anchor(&recipe, &anchor, rng) {
2143 *pos = (*pos + offset + 1) % order.len();
2144 return (Some((recipe, sample)), attempts);
2145 }
2146 }
2147
2148 (None, attempts)
2149 }
2150
2151 fn next_pair_batch_inner_with_weights(
2152 &mut self,
2153 target_split: SplitLabel,
2154 weights: Option<&HashMap<SourceId, f32>>,
2155 ) -> Result<SampleBatch, SamplerError> {
2156 self.ingest_with_weights_fallback(target_split, weights)?;
2157 self.ensure_split_has_records(target_split)?;
2158 let sources = self.source_order.clone();
2159 if sources.is_empty() {
2160 if self.triplet_recipes.is_empty() {
2161 return Err(SamplerError::Configuration(
2162 "no triplet recipes available".into(),
2163 ));
2164 }
2165 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2166 let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2167 let recipe_order = self.recipe_order_weighted_cycled(
2168 &recipe_weights,
2169 self.triplet_recipe_rr_idx,
2170 &mut rng,
2171 );
2172 let mut pairs = Vec::new();
2173 let mut seen = HashSet::new();
2174 let mut last_recipe_name = None;
2175 let mut recipe_pos = 0usize;
2176 let mut recipe_steps = 0usize;
2177 let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2178 for _ in 0..attempts {
2179 if pairs.len() >= self.config.batch_size {
2180 break;
2181 }
2182 let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2183 break;
2184 };
2185 let mut triplet = None;
2186 for offset in 0..recipe_order.len() {
2187 let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2188 recipe_steps = recipe_steps.saturating_add(1);
2189 let recipe = self.triplet_recipes[idx].clone();
2190 last_recipe_name = Some(recipe.name.clone());
2191 if let Some(sample) = self.make_triplet_with_anchor(&recipe, &anchor, &mut rng)
2192 {
2193 triplet = Some((recipe, sample));
2194 recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2195 break;
2196 }
2197 }
2198 if let Some((recipe, triplet)) = triplet {
2199 let key = (
2200 triplet.anchor.record_id.clone(),
2201 triplet.positive.record_id.clone(),
2202 triplet.negative.record_id.clone(),
2203 );
2204 if seen.insert(key) {
2205 let SampleTriplet {
2206 recipe: triplet_recipe_name,
2207 anchor,
2208 positive,
2209 negative,
2210 weight,
2211 instruction,
2212 } = triplet;
2213
2214 let record_label = self
2217 .records
2218 .get(&anchor.record_id)
2219 .and_then(|r| r.label.clone());
2220
2221 if pairs.len() < self.config.batch_size {
2222 pairs.push(SamplePair {
2223 recipe: triplet_recipe_name.clone(),
2224 anchor: anchor.clone(),
2225 positive: positive.clone(),
2226 weight,
2227 instruction: instruction.clone(),
2228 label: record_label.clone().unwrap_or(PairLabel::Positive),
2229 reason: None,
2230 });
2231 }
2232 if pairs.len() < self.config.batch_size {
2233 pairs.push(SamplePair {
2234 recipe: triplet_recipe_name,
2235 anchor,
2236 positive: negative,
2237 weight,
2238 instruction,
2239 label: record_label.unwrap_or(PairLabel::Negative),
2240 reason: Some(
2241 strategy_reason(&recipe.negative_strategy).to_string(),
2242 ),
2243 });
2244 }
2245 }
2246 }
2247 }
2248 if recipe_steps > 0 {
2249 self.triplet_recipe_rr_idx =
2250 self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2251 }
2252 self.rng = rng;
2253 pad_with_reuse(&mut pairs, self.config.batch_size);
2254 if pairs.len() == self.config.batch_size {
2255 return Ok(SampleBatch { pairs });
2256 }
2257 return Err(SamplerError::Exhausted(
2258 last_recipe_name
2259 .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2260 .to_string(),
2261 ));
2262 }
2263
2264 let mut pairs = Vec::new();
2265 let mut seen = HashSet::new();
2266 let mut source_steps = 0usize;
2267 let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2268 let mut source_idx = self.source_cycle_idx % sources.len();
2269 let mut cycle_sources = self.shuffled_source_cycle(cycle);
2270 let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2271 let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2272 let mut recipe_steps = 0usize;
2273 let max_recipe_len = sources
2274 .iter()
2275 .map(|source| self.triplet_recipe_count_for_source(source))
2276 .max()
2277 .unwrap_or(1)
2278 .max(1);
2279 let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2280 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2281 for _ in 0..attempts {
2282 if pairs.len() >= self.config.batch_size {
2283 break;
2284 }
2285 let source = cycle_sources[source_idx].as_str();
2286 let (triplet, attempts_used) = self.sample_source_triplet_candidate(
2287 source,
2288 target_split,
2289 &mut recipe_orders,
2290 &mut recipe_positions,
2291 &mut rng,
2292 );
2293 recipe_steps = recipe_steps.saturating_add(attempts_used);
2294 if let Some((recipe, triplet)) = triplet {
2295 let key = (
2296 triplet.anchor.record_id.clone(),
2297 triplet.positive.record_id.clone(),
2298 triplet.negative.record_id.clone(),
2299 );
2300 if seen.insert(key) {
2301 let SampleTriplet {
2302 recipe: triplet_recipe_name,
2303 anchor,
2304 positive,
2305 negative,
2306 weight,
2307 instruction,
2308 } = triplet;
2309
2310 let record_label = self
2312 .records
2313 .get(&anchor.record_id)
2314 .and_then(|r| r.label.clone());
2315
2316 if pairs.len() < self.config.batch_size {
2317 pairs.push(SamplePair {
2318 recipe: triplet_recipe_name.clone(),
2319 anchor: anchor.clone(),
2320 positive: positive.clone(),
2321 weight,
2322 instruction: instruction.clone(),
2323 label: record_label.clone().unwrap_or(PairLabel::Positive),
2324 reason: None,
2325 });
2326 }
2327 if pairs.len() < self.config.batch_size {
2328 pairs.push(SamplePair {
2329 recipe: triplet_recipe_name,
2330 anchor,
2331 positive: negative,
2332 weight,
2333 instruction,
2334 label: record_label.unwrap_or(PairLabel::Negative),
2335 reason: Some(strategy_reason(&recipe.negative_strategy).to_string()),
2336 });
2337 }
2338 }
2339 }
2340 source_idx += 1;
2341 source_steps += 1;
2342 if source_idx >= cycle_sources.len() {
2343 source_idx = 0;
2344 cycle = cycle.saturating_add(1);
2345 cycle_sources = self.shuffled_source_cycle(cycle);
2346 }
2347 }
2348 if recipe_steps > 0 {
2349 self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2350 }
2351 self.rng = rng;
2352 pad_with_reuse(&mut pairs, self.config.batch_size);
2353 if pairs.len() == self.config.batch_size {
2354 self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2355 self.source_state_dirty = sources.len() > 1;
2356 return Ok(SampleBatch { pairs });
2357 }
2358 Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
2359 }
2360
2361 fn try_emit_text_sample(
2366 &mut self,
2367 sample: &TextSample,
2368 batch_texts: &mut HashSet<String>,
2369 ) -> bool {
2370 batch_texts.insert(sample.chunk.text.clone())
2371 && self
2372 .emitted_text_hashes
2373 .entry(sample.chunk.record_id.clone())
2374 .or_default()
2375 .insert(stable_hash_str(0, &sample.chunk.text))
2376 }
2377
2378 fn next_text_batch_inner_with_weights(
2379 &mut self,
2380 target_split: SplitLabel,
2381 weights: Option<&HashMap<SourceId, f32>>,
2382 ) -> Result<TextBatch, SamplerError> {
2383 self.ingest_with_weights_fallback(target_split, weights)?;
2384 self.ensure_split_has_records(target_split)?;
2385 let sources = self.source_order.clone();
2386 if sources.is_empty() {
2387 if self.text_recipes.is_empty() {
2388 return Err(SamplerError::Configuration(
2389 "no text recipes configured".into(),
2390 ));
2391 }
2392 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2393 let recipe_weights: Vec<f32> = self.text_recipes.iter().map(|r| r.weight).collect();
2394 let recipe_order = self.text_recipe_order_weighted_cycled(
2395 &recipe_weights,
2396 self.text_recipe_rr_idx,
2397 &mut rng,
2398 );
2399 let mut samples = Vec::new();
2400 let mut batch_texts = HashSet::new();
2401 let mut last_recipe_name = None;
2402 let mut recipe_pos = 0usize;
2403 let mut recipe_steps = 0usize;
2404 let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2405 for _ in 0..attempts {
2406 if samples.len() >= self.config.batch_size {
2407 break;
2408 }
2409 let recipe_idx = recipe_order[recipe_pos];
2410 recipe_pos = (recipe_pos + 1) % recipe_order.len();
2411 recipe_steps = recipe_steps.saturating_add(1);
2412 let recipe = self.text_recipes[recipe_idx].clone();
2413 last_recipe_name = Some(recipe.name.clone());
2414 if let Some(sample) =
2415 self.make_text_sample_for_split(&recipe, None, target_split, &mut rng)
2416 && self.try_emit_text_sample(&sample, &mut batch_texts)
2417 {
2418 samples.push(sample);
2419 }
2420 }
2421 if recipe_steps > 0 {
2422 self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2423 }
2424 self.rng = rng;
2425 pad_with_reuse(&mut samples, self.config.batch_size);
2426 if samples.len() == self.config.batch_size {
2427 return Ok(TextBatch { samples });
2428 }
2429 return Err(SamplerError::Exhausted(
2430 last_recipe_name
2431 .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TEXT))
2432 .to_string(),
2433 ));
2434 }
2435
2436 let mut samples = Vec::new();
2437 let mut batch_texts = HashSet::new();
2438 let mut source_steps = 0usize;
2439 let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2440 let mut idx = self.source_cycle_idx % sources.len();
2441 let mut cycle_sources = self.shuffled_source_cycle(cycle);
2442 let mut recipe_orders: HashMap<RecipeKey, Vec<usize>> = HashMap::new();
2443 let mut recipe_positions: HashMap<RecipeKey, usize> = HashMap::new();
2444 let mut recipe_steps = 0usize;
2445 let max_recipe_len = sources
2446 .iter()
2447 .map(|source| self.text_recipes_for_source(source).len())
2448 .max()
2449 .unwrap_or(1)
2450 .max(1);
2451 let attempts = self.config.batch_size * 4 * sources.len() * max_recipe_len;
2452 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2453 for _ in 0..attempts {
2454 if samples.len() >= self.config.batch_size {
2455 break;
2456 }
2457 let source = cycle_sources[idx].as_str();
2458 let recipes = self.text_recipes_for_source(source).to_vec();
2459 if recipes.is_empty() {
2460 idx += 1;
2461 source_steps += 1;
2462 if idx >= cycle_sources.len() {
2463 idx = 0;
2464 cycle = cycle.saturating_add(1);
2465 cycle_sources = self.shuffled_source_cycle(cycle);
2466 }
2467 continue;
2468 }
2469 if !recipe_orders.contains_key(source) {
2470 let recipe_weights: Vec<f32> = recipes.iter().map(|r| r.weight).collect();
2471 let order = self.text_recipe_order_weighted_cycled(
2472 &recipe_weights,
2473 self.text_recipe_rr_idx,
2474 &mut rng,
2475 );
2476 recipe_orders.insert(source.to_string(), order);
2477 }
2478 let order = recipe_orders
2479 .get(source)
2480 .expect("recipe order missing for source");
2481 let pos = recipe_positions.entry(source.to_string()).or_insert(0);
2482 let mut sample: Option<(TextRecipe, TextSample)> = None;
2483 for offset in 0..order.len() {
2484 let recipe_idx = order[(*pos + offset) % order.len()];
2485 let recipe = recipes[recipe_idx].clone();
2486 if let Some(item) =
2487 self.make_text_sample_for_split(&recipe, Some(source), target_split, &mut rng)
2488 && self.try_emit_text_sample(&item, &mut batch_texts)
2489 {
2490 recipe_steps = recipe_steps.saturating_add(offset + 1);
2491 *pos = (*pos + offset + 1) % order.len();
2492 sample = Some((recipe, item));
2493 break;
2494 }
2495 }
2496 if sample.is_none() {
2497 recipe_steps = recipe_steps.saturating_add(order.len());
2498 }
2499 if let Some((_recipe, sample)) = sample {
2500 samples.push(sample);
2501 }
2502 idx += 1;
2503 source_steps += 1;
2504 if idx >= cycle_sources.len() {
2505 idx = 0;
2506 cycle = cycle.saturating_add(1);
2507 cycle_sources = self.shuffled_source_cycle(cycle);
2508 }
2509 }
2510 if samples.len() != self.config.batch_size {
2511 pad_with_reuse(&mut samples, self.config.batch_size);
2512 }
2513 if samples.len() != self.config.batch_size {
2514 self.rng = rng;
2515 return Err(SamplerError::Exhausted(RECIPE_LABEL_TEXT.into()));
2516 }
2517 self.rng = rng;
2518 self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2519 self.source_state_dirty = sources.len() > 1;
2520 if recipe_steps > 0 {
2521 self.text_recipe_rr_idx = self.text_recipe_rr_idx.saturating_add(recipe_steps);
2522 }
2523 Ok(TextBatch { samples })
2524 }
2525
2526 fn next_triplet_batch_inner_with_weights(
2527 &mut self,
2528 target_split: SplitLabel,
2529 weights: Option<&HashMap<SourceId, f32>>,
2530 ) -> Result<TripletBatch, SamplerError> {
2531 self.ingest_with_weights_fallback(target_split, weights)?;
2532 self.ensure_split_has_records(target_split)?;
2533 let sources = self.source_order.clone();
2534 if sources.is_empty() {
2535 if self.triplet_recipes.is_empty() {
2536 return Err(SamplerError::Configuration(
2537 "no triplet recipes configured".into(),
2538 ));
2539 }
2540 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2541 let recipe_weights: Vec<f32> = self.triplet_recipes.iter().map(|r| r.weight).collect();
2542 let recipe_order = self.recipe_order_weighted_cycled(
2543 &recipe_weights,
2544 self.triplet_recipe_rr_idx,
2545 &mut rng,
2546 );
2547 let mut triplets = Vec::new();
2548 let mut seen = HashSet::new();
2549 let mut last_recipe_name = None;
2550 let mut recipe_pos = 0usize;
2551 let mut recipe_steps = 0usize;
2552 let attempts = self.config.batch_size * 4 * recipe_order.len().max(1);
2553 for _ in 0..attempts {
2554 if triplets.len() >= self.config.batch_size {
2555 break;
2556 }
2557 let Some(anchor) = self.choose_anchor_record(None, target_split) else {
2558 break;
2559 };
2560 let mut triplet = None;
2561 for offset in 0..recipe_order.len() {
2562 let idx = recipe_order[(recipe_pos + offset) % recipe_order.len()];
2563 recipe_steps = recipe_steps.saturating_add(1);
2564 let recipe = self.triplet_recipes[idx].clone();
2565 last_recipe_name = Some(recipe.name.clone());
2566 if let Some(sample) = self.make_triplet_with_anchor(&recipe, &anchor, &mut rng)
2567 {
2568 triplet = Some(sample);
2569 recipe_pos = (recipe_pos + offset + 1) % recipe_order.len();
2570 break;
2571 }
2572 }
2573 if let Some(triplet) = triplet {
2574 let key = (
2575 triplet.anchor.record_id.clone(),
2576 triplet.positive.record_id.clone(),
2577 triplet.negative.record_id.clone(),
2578 );
2579 if seen.insert(key) {
2580 triplets.push(triplet);
2581 }
2582 }
2583 }
2584 if recipe_steps > 0 {
2585 self.triplet_recipe_rr_idx =
2586 self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2587 }
2588 self.rng = rng;
2589 pad_with_reuse(&mut triplets, self.config.batch_size);
2590 if triplets.len() == self.config.batch_size {
2591 return Ok(TripletBatch { triplets });
2592 }
2593 return Err(SamplerError::Exhausted(
2594 last_recipe_name
2595 .unwrap_or(Cow::Borrowed(RECIPE_LABEL_TRIPLETS))
2596 .to_string(),
2597 ));
2598 }
2599
2600 let mut triplets: Vec<SampleTriplet> = Vec::new();
2602 let mut seen: HashSet<(RecordId, RecordId, RecordId)> = HashSet::new();
2603 let mut source_steps = 0usize;
2604 let mut cycle = (self.source_cycle_idx / sources.len()) as u64;
2605 let mut source_idx = self.source_cycle_idx % sources.len();
2606 let mut cycle_sources = self.shuffled_source_cycle(cycle);
2607 let mut recipe_steps = 0usize;
2608 let max_recipe_len = sources
2609 .iter()
2610 .map(|source| self.triplet_recipe_count_for_source(source.as_str()))
2611 .max()
2612 .unwrap_or(1)
2613 .max(1);
2614 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2615
2616 struct SlotPlan {
2621 anchor: Arc<DataRecord>,
2622 recipes: Vec<TripletRecipe>,
2623 fork_seed: u64,
2624 }
2625 let target_slots = self.config.batch_size * 4 * max_recipe_len;
2626 let mut slot_plans: Vec<SlotPlan> = Vec::with_capacity(target_slots);
2627
2628 for _ in 0..target_slots {
2629 if slot_plans.len() >= target_slots {
2630 break;
2631 }
2632 let source = cycle_sources[source_idx].as_str();
2633 let (recipes, _) = self.resolve_source_triplet_plan(source);
2634 if !recipes.is_empty() {
2635 let fork_seed = rng.next_u64();
2636 if let Some(anchor) = self.choose_anchor_record(Some(source), target_split) {
2637 slot_plans.push(SlotPlan {
2638 anchor,
2639 recipes,
2640 fork_seed,
2641 });
2642 recipe_steps = recipe_steps.saturating_add(1);
2643 }
2644 }
2645 source_idx += 1;
2646 source_steps += 1;
2647 if source_idx >= cycle_sources.len() {
2648 source_idx = 0;
2649 cycle = cycle.saturating_add(1);
2650 cycle_sources = self.shuffled_source_cycle(cycle);
2651 }
2652 }
2653 self.rng = rng;
2654
2655 struct SlotCandidate {
2660 recipe: TripletRecipe,
2661 anchor: Arc<DataRecord>,
2662 anchor_chunk: RecordChunk,
2663 positive_chunk: RecordChunk,
2664 anchor_raw_text: String,
2665 }
2666 let mut raw_candidates: Vec<(usize, Option<SlotCandidate>)> = slot_plans
2667 .par_iter()
2668 .enumerate()
2669 .map(|(slot_idx, plan)| {
2670 let mut fork_rng = DeterministicRng::new(plan.fork_seed);
2671 let weights: Vec<f32> = plan.recipes.iter().map(|r| r.weight).collect();
2674 let order = weighted_recipe_order(&weights, &mut fork_rng);
2675 let mut candidate = None;
2676 for &idx in &order {
2677 let recipe = &plan.recipes[idx];
2678 if let Some((ac, pc, raw)) =
2679 self.select_anchor_positive_for_recipe(recipe, &plan.anchor, &mut fork_rng)
2680 {
2681 candidate = Some(SlotCandidate {
2682 recipe: recipe.clone(),
2683 anchor: Arc::clone(&plan.anchor),
2684 anchor_chunk: ac,
2685 positive_chunk: pc,
2686 anchor_raw_text: raw,
2687 });
2688 break;
2689 }
2690 }
2691 (slot_idx, candidate)
2692 })
2693 .collect();
2694
2695 raw_candidates.sort_unstable_by_key(|(i, _)| *i);
2700 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2701 for (_, candidate) in raw_candidates {
2702 if triplets.len() >= self.config.batch_size {
2703 break;
2704 }
2705 let Some(sc) = candidate else { continue };
2706 let (negative_record, fallback_used) = match self.select_negative_record(
2707 &sc.anchor,
2708 &sc.recipe.negative_strategy,
2709 Some(&sc.anchor_raw_text),
2710 &mut rng,
2711 ) {
2712 Some(pair) => pair,
2713 None => continue,
2714 };
2715 let mut negative_chunk = match self.select_chunk_parallel(
2716 &negative_record,
2717 &sc.recipe.negative_selector,
2718 &mut rng,
2719 ) {
2720 Some(c) => c,
2721 None => continue,
2722 };
2723 self.decorate_chunk_parallel(&negative_record, &mut negative_chunk, &mut rng);
2724
2725 let (anchor_chunk, positive_chunk) = if rng.next_u64() & ANCHOR_POSITIVE_SWAP_MASK == 0
2727 {
2728 (sc.positive_chunk, sc.anchor_chunk)
2729 } else {
2730 (sc.anchor_chunk, sc.positive_chunk)
2731 };
2732
2733 if (!sc.recipe.allow_same_anchor_positive && anchor_chunk.text == positive_chunk.text)
2734 || negative_chunk.text == positive_chunk.text
2735 || negative_chunk.text == anchor_chunk.text
2736 {
2737 continue;
2738 }
2739
2740 let chunk_weight = self.triplet_chunk_weight(
2741 &sc.recipe,
2742 &anchor_chunk,
2743 &positive_chunk,
2744 &negative_chunk,
2745 );
2746 let weight = sc.recipe.weight * chunk_weight;
2747 let recipe_name = if fallback_used {
2748 format!("{}_fallback_same_split", sc.recipe.name)
2749 } else {
2750 sc.recipe.name.to_string()
2751 };
2752 let triplet = SampleTriplet {
2753 recipe: recipe_name,
2754 anchor: anchor_chunk,
2755 positive: positive_chunk,
2756 negative: negative_chunk,
2757 weight,
2758 instruction: sc.recipe.instruction.as_ref().map(|s| s.to_string()),
2759 };
2760 let key = (
2761 triplet.anchor.record_id.clone(),
2762 triplet.positive.record_id.clone(),
2763 triplet.negative.record_id.clone(),
2764 );
2765 if seen.insert(key) && triplets.len() < self.config.batch_size {
2766 triplets.push(triplet);
2767 }
2768 }
2769 self.rng = rng;
2770
2771 if recipe_steps > 0 {
2772 self.triplet_recipe_rr_idx = self.triplet_recipe_rr_idx.saturating_add(recipe_steps);
2773 }
2774 pad_with_reuse(&mut triplets, self.config.batch_size);
2775 if triplets.len() == self.config.batch_size {
2776 self.source_cycle_idx = self.source_cycle_idx.saturating_add(source_steps);
2777 self.source_state_dirty = sources.len() > 1;
2778 let batch = TripletBatch { triplets };
2779 return Ok(batch);
2780 }
2781 Err(SamplerError::Exhausted(RECIPE_LABEL_TRIPLETS.into()))
2782 }
2783
2784 #[cfg(test)]
2790 #[cfg(feature = "bm25-mining")]
2791 fn bm25_backend_mut(&mut self) -> &mut backends::Bm25Backend {
2792 self.negative_backend
2793 .as_any_mut()
2794 .downcast_mut::<backends::Bm25Backend>()
2795 .expect("bm25_backend_mut: negative_backend is Bm25Backend when bm25-mining feature is active")
2796 }
2797
2798 #[cfg(test)]
2800 fn recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
2801 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2802 let result = self.recipe_order_weighted_shuffled(weights, &mut rng);
2803 self.rng = rng;
2804 result
2805 }
2806
2807 #[cfg(test)]
2809 fn recipe_order_weighted_cycled_seeded(
2810 &mut self,
2811 weights: &[f32],
2812 rr_idx: usize,
2813 ) -> Vec<usize> {
2814 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2815 let result = self.recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
2816 self.rng = rng;
2817 result
2818 }
2819
2820 #[cfg(test)]
2822 fn text_recipe_order_weighted_shuffled_seeded(&mut self, weights: &[f32]) -> Vec<usize> {
2823 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2824 let result = self.text_recipe_order_weighted_shuffled(weights, &mut rng);
2825 self.rng = rng;
2826 result
2827 }
2828
2829 #[cfg(test)]
2831 fn text_recipe_order_weighted_cycled_seeded(
2832 &mut self,
2833 weights: &[f32],
2834 rr_idx: usize,
2835 ) -> Vec<usize> {
2836 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2837 let result = self.text_recipe_order_weighted_cycled(weights, rr_idx, &mut rng);
2838 self.rng = rng;
2839 result
2840 }
2841
2842 #[cfg(test)]
2844 fn select_negative_record_seeded(
2845 &mut self,
2846 anchor_record: &DataRecord,
2847 strategy: &NegativeStrategy,
2848 anchor_query_text: Option<&str>,
2849 ) -> Option<(Arc<DataRecord>, bool)> {
2850 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2851 let result =
2852 self.select_negative_record(anchor_record, strategy, anchor_query_text, &mut rng);
2853 self.rng = rng;
2854 result
2855 }
2856
2857 #[cfg(test)]
2859 fn make_triplet_with_anchor_seeded(
2860 &mut self,
2861 recipe: &TripletRecipe,
2862 anchor: &DataRecord,
2863 ) -> Option<SampleTriplet> {
2864 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2865 let result = self.make_triplet_with_anchor(recipe, anchor, &mut rng);
2866 self.rng = rng;
2867 result
2868 }
2869
2870 #[cfg(test)]
2872 fn decorate_chunk_seeded(&mut self, record: &DataRecord, chunk: &mut RecordChunk) {
2873 let mut rng = std::mem::replace(&mut self.rng, DeterministicRng::new(0));
2874 self.decorate_chunk(record, chunk, &mut rng);
2875 self.rng = rng;
2876 }
2877
2878 #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
2884 fn bm25_fallback_stats(&self) -> (u64, u64) {
2885 self.negative_backend.bm25_fallback_stats()
2886 }
2887
2888 #[cfg(test)]
2893 #[cfg(feature = "bm25-mining")]
2894 fn bm25_ranked_candidates(&mut self, anchor: &crate::data::DataRecord) -> Vec<RecordId> {
2895 let split = self
2896 .split_store
2897 .label_for(&anchor.id)
2898 .unwrap_or(SplitLabel::Train);
2899 self.negative_backend
2900 .as_any_mut()
2901 .downcast_mut::<backends::Bm25Backend>()
2902 .expect("bm25_ranked_candidates: Bm25Backend")
2903 .ranked_candidates_pub(anchor, split)
2904 }
2905}
2906
2907fn weighted_recipe_order(weights: &[f32], rng: &mut DeterministicRng) -> Vec<usize> {
2918 let nonzero: Vec<(usize, f32)> = weights
2919 .iter()
2920 .enumerate()
2921 .filter(|(_, w)| **w > 0.0)
2922 .map(|(i, &w)| (i, w))
2923 .collect();
2924 if nonzero.is_empty() {
2925 return Vec::new();
2926 }
2927 let w_min = nonzero
2928 .iter()
2929 .map(|(_, w)| *w)
2930 .fold(f32::INFINITY, f32::min);
2931 let mut order: Vec<usize> = Vec::new();
2932 for (recipe_idx, w) in &nonzero {
2933 let tickets = ((w / w_min).round() as usize).clamp(1, RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER);
2934 for _ in 0..tickets {
2935 order.push(*recipe_idx);
2936 }
2937 }
2938 order.shuffle(rng);
2939 order
2940}
2941
2942fn same_selector_pair_is_valid(
2943 anchor_chunk: &RecordChunk,
2944 positive_chunk: &RecordChunk,
2945 enforce_window_pair: bool,
2946) -> bool {
2947 if triplet_chunk_key(anchor_chunk) == triplet_chunk_key(positive_chunk) {
2948 return false;
2949 }
2950 if !enforce_window_pair {
2951 return true;
2952 }
2953 matches!(
2954 (&anchor_chunk.view, &positive_chunk.view),
2955 (ChunkView::Window { .. }, ChunkView::Window { .. })
2956 )
2957}
2958
2959impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> TripletSampler<S> {
2960 pub fn new(config: SamplerConfig, split_store: Arc<S>) -> Self {
2962 let inner = TripletSamplerInner::new(config, split_store);
2963 Self {
2964 inner: Mutex::new(inner),
2965 }
2966 }
2967
2968 pub fn new_with_chunker(
2970 config: SamplerConfig,
2971 split_store: Arc<S>,
2972 chunker: Arc<dyn ChunkingAlgorithm>,
2973 ) -> Self {
2974 let inner = TripletSamplerInner::new_with_chunker(config, split_store, chunker);
2975 Self {
2976 inner: Mutex::new(inner),
2977 }
2978 }
2979
2980 #[deprecated(
2986 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"
2987 )]
2988 pub fn next_pair_batch_for_split(
2989 &self,
2990 split: SplitLabel,
2991 ) -> Result<SampleBatch, SamplerError> {
2992 self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
2993 }
2994
2995 #[deprecated(
3001 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"
3002 )]
3003 pub fn next_text_batch_for_split(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3004 self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3005 }
3006
3007 #[deprecated(
3013 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"
3014 )]
3015 pub fn next_triplet_batch_for_split(
3016 &self,
3017 split: SplitLabel,
3018 ) -> Result<TripletBatch, SamplerError> {
3019 self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3020 }
3021
3022 fn next_batch_with_retry<B>(
3026 &self,
3027 split: SplitLabel,
3028 weights: &HashMap<SourceId, f32>,
3029 inner_fn: impl Fn(
3030 &mut TripletSamplerInner<S>,
3031 SplitLabel,
3032 Option<&HashMap<SourceId, f32>>,
3033 ) -> Result<B, SamplerError>,
3034 ) -> Result<B, SamplerError> {
3035 let mut inner = self.inner.lock().unwrap();
3036 inner.ensure_split_allowed(split)?;
3037 inner.ensure_ingestion_cursors_loaded()?;
3040 inner.ingestion.increment_epoch_step();
3042 for attempt in 0..=EXHAUSTION_RETRY_LIMIT {
3043 match inner_fn(&mut inner, split, Some(weights)) {
3044 Ok(batch) => return Ok(batch),
3045 Err(SamplerError::Exhausted(_)) if attempt < EXHAUSTION_RETRY_LIMIT => {
3046 inner.force_ingest_refresh_with_weights_for_split(split, weights)?;
3047 }
3048 Err(err) => return Err(err),
3049 }
3050 }
3051 unreachable!()
3053 }
3054
3055 pub fn next_pair_batch_with_weights_for_split(
3057 &self,
3058 split: SplitLabel,
3059 weights: &HashMap<SourceId, f32>,
3060 ) -> Result<SampleBatch, SamplerError> {
3061 self.next_batch_with_retry(split, weights, |inner, split, weights| {
3062 inner.next_pair_batch_inner_with_weights(split, weights)
3063 })
3064 }
3065
3066 pub fn next_text_batch_with_weights_for_split(
3068 &self,
3069 split: SplitLabel,
3070 weights: &HashMap<SourceId, f32>,
3071 ) -> Result<TextBatch, SamplerError> {
3072 self.next_batch_with_retry(split, weights, |inner, split, weights| {
3073 inner.next_text_batch_inner_with_weights(split, weights)
3074 })
3075 }
3076
3077 pub fn next_triplet_batch_with_weights_for_split(
3079 &self,
3080 split: SplitLabel,
3081 weights: &HashMap<SourceId, f32>,
3082 ) -> Result<TripletBatch, SamplerError> {
3083 self.next_batch_with_retry(split, weights, |inner, split, weights| {
3084 inner.next_triplet_batch_inner_with_weights(split, weights)
3085 })
3086 }
3087
3088 #[deprecated(
3094 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"
3095 )]
3096 pub fn prefetch_triplet_batches(
3097 self: Arc<Self>,
3098 split: SplitLabel,
3099 capacity: usize,
3100 ) -> BatchPrefetcher<TripletBatch> {
3101 BatchPrefetcher::new(capacity, move || {
3102 self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3103 })
3104 }
3105
3106 pub fn prefetch_triplet_batches_with_weights(
3108 self: Arc<Self>,
3109 split: SplitLabel,
3110 capacity: usize,
3111 weights: HashMap<SourceId, f32>,
3112 ) -> BatchPrefetcher<TripletBatch> {
3113 BatchPrefetcher::new(capacity, move || {
3114 self.next_triplet_batch_with_weights_for_split(split, &weights)
3115 })
3116 }
3117
3118 #[deprecated(
3124 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"
3125 )]
3126 pub fn prefetch_pair_batches(
3127 self: Arc<Self>,
3128 split: SplitLabel,
3129 capacity: usize,
3130 ) -> BatchPrefetcher<SampleBatch> {
3131 BatchPrefetcher::new(capacity, move || {
3132 self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3133 })
3134 }
3135
3136 pub fn prefetch_pair_batches_with_weights(
3138 self: Arc<Self>,
3139 split: SplitLabel,
3140 capacity: usize,
3141 weights: HashMap<SourceId, f32>,
3142 ) -> BatchPrefetcher<SampleBatch> {
3143 BatchPrefetcher::new(capacity, move || {
3144 self.next_pair_batch_with_weights_for_split(split, &weights)
3145 })
3146 }
3147
3148 #[deprecated(
3154 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"
3155 )]
3156 pub fn prefetch_text_batches(
3157 self: Arc<Self>,
3158 split: SplitLabel,
3159 capacity: usize,
3160 ) -> BatchPrefetcher<TextBatch> {
3161 BatchPrefetcher::new(capacity, move || {
3162 self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3163 })
3164 }
3165
3166 pub fn prefetch_text_batches_with_weights(
3168 self: Arc<Self>,
3169 split: SplitLabel,
3170 capacity: usize,
3171 weights: HashMap<SourceId, f32>,
3172 ) -> BatchPrefetcher<TextBatch> {
3173 BatchPrefetcher::new(capacity, move || {
3174 self.next_text_batch_with_weights_for_split(split, &weights)
3175 })
3176 }
3177
3178 pub fn text_recipes(&self) -> Vec<TextRecipe> {
3180 let inner = self.inner.lock().unwrap();
3181 inner.text_recipes().to_vec()
3182 }
3183
3184 pub fn register_source(
3189 &self,
3190 source: Box<dyn DataSource + 'static>,
3191 ) -> Result<(), SamplerError> {
3192 let mut inner = self.inner.lock().unwrap();
3193 inner.register_source(source)
3194 }
3195
3196 pub fn set_epoch(&self, epoch: u64) -> Result<(), SamplerError> {
3198 let mut inner = self.inner.lock().unwrap();
3199 inner.set_epoch(epoch)
3200 }
3201
3202 pub fn save_sampler_state(&self, save_to: Option<&Path>) -> Result<(), SamplerError> {
3207 let mut inner = self.inner.lock().unwrap();
3208 inner.save_sampler_state(save_to)
3209 }
3210
3211 #[cfg(all(feature = "bm25-mining", feature = "extended-metrics"))]
3219 pub fn bm25_fallback_stats(&self) -> (u64, u64) {
3220 let inner = self.inner.lock().unwrap();
3221 inner.bm25_fallback_stats()
3222 }
3223}
3224
3225impl<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> Sampler for TripletSampler<S> {
3226 fn next_pair_batch(&self, split: SplitLabel) -> Result<SampleBatch, SamplerError> {
3227 self.next_pair_batch_with_weights_for_split(split, &HashMap::new())
3228 }
3229
3230 fn next_pair_batch_with_weights(
3231 &self,
3232 split: SplitLabel,
3233 weights: &HashMap<SourceId, f32>,
3234 ) -> Result<SampleBatch, SamplerError> {
3235 self.next_pair_batch_with_weights_for_split(split, weights)
3236 }
3237
3238 fn next_text_batch(&self, split: SplitLabel) -> Result<TextBatch, SamplerError> {
3239 self.next_text_batch_with_weights_for_split(split, &HashMap::new())
3240 }
3241
3242 fn next_text_batch_with_weights(
3243 &self,
3244 split: SplitLabel,
3245 weights: &HashMap<SourceId, f32>,
3246 ) -> Result<TextBatch, SamplerError> {
3247 self.next_text_batch_with_weights_for_split(split, weights)
3248 }
3249
3250 fn next_triplet_batch(&self, split: SplitLabel) -> Result<TripletBatch, SamplerError> {
3251 self.next_triplet_batch_with_weights_for_split(split, &HashMap::new())
3252 }
3253
3254 fn next_triplet_batch_with_weights(
3255 &self,
3256 split: SplitLabel,
3257 weights: &HashMap<SourceId, f32>,
3258 ) -> Result<TripletBatch, SamplerError> {
3259 self.next_triplet_batch_with_weights_for_split(split, weights)
3260 }
3261
3262 fn save_sampler_state(&self, save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
3263 TripletSampler::save_sampler_state(self, save_to)
3264 }
3265}
3266
3267fn roles_match(target: &SectionRole, candidate: &SectionRole) -> bool {
3268 target == candidate
3269}
3270
3271fn role_cursor_key(record_id: &RecordId, role: &SectionRole) -> (RecordId, String) {
3272 (record_id.clone(), role_label(role))
3273}
3274
3275fn role_label(role: &SectionRole) -> String {
3276 match role {
3277 SectionRole::Anchor => ROLE_LABEL_ANCHOR.into(),
3278 SectionRole::Context => ROLE_LABEL_CONTEXT.into(),
3279 }
3280}
3281
3282fn taxonomy_value(record: &DataRecord, field: MetadataKey) -> Option<&str> {
3283 record.taxonomy.iter().find_map(|entry| field.strip(entry))
3284}
3285
3286fn strategy_reason(strategy: &NegativeStrategy) -> &'static str {
3287 match strategy {
3288 NegativeStrategy::WrongPublicationDate => NEG_REASON_WRONG_DATE,
3289 NegativeStrategy::WrongArticle => NEG_REASON_WRONG_ARTICLE,
3290 NegativeStrategy::QuestionAnswerMismatch => NEG_REASON_WRONG_QA,
3291 NegativeStrategy::SameRecord => NEG_REASON_SAME_RECORD,
3292 }
3293}
3294
3295#[cfg(test)]
3303fn text_dedup_key(chunk: &RecordChunk) -> (String, String) {
3304 (chunk.record_id.clone(), chunk.text.clone())
3305}
3306
3307fn triplet_chunk_key(chunk: &RecordChunk) -> String {
3308 match &chunk.view {
3309 ChunkView::Window { index, .. } => {
3310 format!("{}|{}|w|{}", chunk.record_id, chunk.section_idx, index)
3311 }
3312 ChunkView::SummaryFallback { strategy, .. } => {
3313 format!("{}|{}|s|{}", chunk.record_id, chunk.section_idx, strategy)
3314 }
3315 }
3316}
3317
3318fn pad_with_reuse<T: Clone>(items: &mut Vec<T>, target: usize) {
3319 if items.is_empty() || items.len() >= target {
3320 return;
3321 }
3322 let seed = items.clone();
3323 let base_len = seed.len();
3324 for idx in 0..(target - items.len()) {
3325 items.push(seed[idx % base_len].clone());
3326 }
3327}
3328
3329#[cfg(test)]
3330mod tests;