Skip to main content

triplets_core/sampler/
mod.rs

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// AUTO-RECIPE HANDLING OVERVIEW (end-to-end):
46// Stage A: Source-level injection eligibility ("should this source even get the recipe?")
47//   - `sync_records_from_cache` -> `record_has_long_anchor_or_context_section` marks
48//     `sources_with_long_sections`.
49//   - `resolve_source_triplet_plan` -> `should_auto_inject_chunk_pair_recipe` appends
50//     `source_chunk_pair_recipe` for eligible sources.
51// Stage B: Record-level execution eligibility ("can this specific record run it now?")
52//   - `make_auto_chunk_pair_triplet_with_anchor` checks
53//     `record_has_at_least_two_window_chunks_for_selector` before sampling.
54// Stage C: Actual chunk window layout ("how are windows formed?")
55//   - `materialize_chunks` is the source of truth for token window construction.
56
57#[derive(Debug, Clone)]
58/// Small deterministic RNG used for reproducible sampler behavior.
59struct 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
107/// Public helper so callers (and examples) can compute the same chunk weighting the sampler uses.
108///
109/// Definitions:
110/// - `trust` (range 0.0-1.0): comes from `RecordChunk.quality.trust`; scales the contribution of a chunk.
111/// - `base`: window chunks use proximity-to-head (`1 / (index + 1)`); summary chunks use their configured `summary_fallback_weight`.
112/// - `chunk_weight_floor`: minimum weight applied after scaling.
113///
114/// Formula: `max(chunk_weight_floor, base * trust)`.
115pub 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
125/// Public sampling interface for pair, triplet, and text batch generation.
126pub trait Sampler {
127    /// Returns a batch of pairs. Consumes the shared epoch cursor for anchor selection.
128    ///
129    /// Samples all sources uniformly. Prefer [`Sampler::next_pair_batch_with_weights`]
130    /// with an explicit per-source weight map to honor a data mixture.
131    #[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    /// Returns a batch of text samples. Consumes the shared epoch cursor for anchor selection.
138    ///
139    /// Samples all sources uniformly. Prefer [`Sampler::next_text_batch_with_weights`]
140    /// with an explicit per-source weight map to honor a data mixture.
141    #[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    /// Returns a batch of triplets. Consumes the shared epoch cursor for anchor selection.
148    ///
149    /// Samples all sources uniformly. Prefer [`Sampler::next_triplet_batch_with_weights`]
150    /// with an explicit per-source weight map to honor a data mixture.
151    #[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    /// Returns a batch of pairs with per-call source weights.
158    fn next_pair_batch_with_weights(
159        &self,
160        split: SplitLabel,
161        weights: &HashMap<SourceId, f32>,
162    ) -> Result<SampleBatch, SamplerError>;
163    /// Returns a batch of text samples with per-call source weights.
164    fn next_text_batch_with_weights(
165        &self,
166        split: SplitLabel,
167        weights: &HashMap<SourceId, f32>,
168    ) -> Result<TextBatch, SamplerError>;
169    /// Returns a batch of triplets with per-call source weights.
170    fn next_triplet_batch_with_weights(
171        &self,
172        split: SplitLabel,
173        weights: &HashMap<SourceId, f32>,
174    ) -> Result<TripletBatch, SamplerError>;
175    /// Persist sampler and split runtime state for restart-resume.
176    ///
177    /// Default is a no-op, allowing mock/test implementations to skip persistence.
178    fn save_sampler_state(&self, _save_to: Option<&std::path::Path>) -> Result<(), SamplerError> {
179        Ok(())
180    }
181}
182
183/// Background prefetcher that fills a bounded queue with sample batches.
184pub 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)]
191/// Prefetcher runtime counters.
192struct 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    /// Block until the next prefetched batch is available.
227    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    /// Number of prefetched batches currently queued.
251    pub fn queue_len(&self) -> usize {
252        self.stats.queued.load(Ordering::Relaxed)
253    }
254
255    /// Total number of batches produced by the background worker.
256    pub fn produced_count(&self) -> usize {
257        self.stats.produced.load(Ordering::Relaxed)
258    }
259
260    /// Total number of errors produced by the background worker.
261    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
275/// Sampler that draws anchors from a single shared epoch cursor and then
276/// selects chunks from those records. Ingestion happens on demand when sampling.
277pub struct TripletSampler<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
278    inner: Mutex<TripletSamplerInner<S>>,
279}
280
281/// Internal sampler state implementation guarded by `TripletSampler`.
282struct TripletSamplerInner<S: SplitStore + EpochStateStore + SamplerStateStore + 'static> {
283    /// Immutable sampler configuration (seed, batch size, recipes, splits, etc.).
284    config: SamplerConfig,
285    /// Active chunking implementation used to materialize section chunks.
286    chunker: Arc<dyn ChunkingAlgorithm>,
287    /// Split store backing train/val/test assignments and persisted sampler state.
288    split_store: Arc<S>,
289    /// On-demand ingestion manager that fills the batch-sized buffer.
290    ingestion: IngestionManager,
291    /// Current in-memory record pool keyed by record id.
292    records: IndexMap<RecordId, Arc<DataRecord>>,
293    /// Deterministic RNG for per-batch shuffles and sampling.
294    rng: DeterministicRng,
295    /// Config-level triplet recipes used when sources do not supply their own.
296    triplet_recipes: Vec<TripletRecipe>,
297    /// Config-level text recipes used when sources do not supply their own.
298    text_recipes: Vec<TextRecipe>,
299    /// Per-source triplet recipes keyed by source id.
300    source_triplet_recipes: HashMap<SourceId, Vec<TripletRecipe>>,
301    /// Sources that currently contain at least one section larger than the chunk window.
302    sources_with_long_sections: HashSet<SourceId>,
303    /// Per-source text recipes keyed by source id.
304    source_text_recipes: HashMap<SourceId, Vec<TextRecipe>>,
305    /// True if triplet recipes came from config (no source defaults).
306    using_config_triplet_recipes: bool,
307    /// True if text recipes came from config (no derivation).
308    using_config_text_recipes: bool,
309    /// Last observed ingestion counter for cache updates.
310    last_observed_ingest: u64,
311    /// Epoch tracker for split-aware deterministic sampling.
312    epoch_tracker: EpochTracker,
313    /// Per-record, per-section chunk cursor to rotate through chunk windows.
314    chunk_cursors: HashMap<(RecordId, usize), usize>,
315    /// Per-record, per-role cursor to rotate through role-specific sections.
316    role_cursors: HashMap<(RecordId, String), usize>,
317    /// Pluggable negative-selection backend (uniform-random or BM25).
318    negative_backend: Box<dyn backends::NegativeBackend>,
319    /// Chunk id to record id lookup (used by epoch tracker).
320    chunk_index: HashMap<RecordId, RecordId>,
321    /// Round-robin order of source ids (deterministic).
322    source_order: Vec<SourceId>,
323    /// Current index into `source_order` for shuffled-cycle sampling.
324    source_cycle_idx: usize,
325    /// True once persisted source state is loaded.
326    source_state_loaded: bool,
327    /// True once persisted ingestion cursors are loaded.
328    ingestion_cursors_loaded: bool,
329    /// Marks whether source state needs persistence.
330    source_state_dirty: bool,
331    /// Record indices per source for round-robin within a source.
332    source_record_indices: HashMap<SourceId, Vec<usize>>,
333    /// Per-source cursor into `source_record_indices`.
334    source_record_cursors: HashMap<SourceId, usize>,
335    /// Round-robin index for triplet recipe cycling.
336    triplet_recipe_rr_idx: usize,
337    /// Hashes of text already emitted by the text batch path while the record
338    /// pool is stable.  Pruned on sync_records_from_cache only when records
339    /// are evicted.  Prevents the same text from being sampled again within
340    /// the same ingestion window, regardless of source wrapping or epoch.
341    ///
342    /// Per-record tracking lets us surgically remove hashes only for evicted
343    /// records, preserving cross-batch dedup across small window advances
344    /// (e.g. window=512, batch_size=4 where only 4/512 records are swapped).
345    emitted_text_hashes: HashMap<RecordId, HashSet<u64>>,
346    /// Round-robin index for text recipe cycling.
347    text_recipe_rr_idx: usize,
348    /// Epoch counter for per-source deterministic shuffling (seed ^ epoch).
349    epoch: u64,
350
351    /// Tracks whether each source has wrapped its cursor in the current epoch.
352    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    /// Current epoch-adjusted seed: mixes `epoch` into `config.seed` so every epoch
440    /// produces a distinct permutation across all seed-dependent operations.
441    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        // Reset epoch step at epoch boundary so each epoch
474        // starts with step=0, giving deterministic step sequences.
475        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            // First touch should **not** always start at window 0.
498            // We derive a deterministic per-(record,section) starting offset so:
499            // - repeated runs with the same seed/epoch are reproducible,
500            // - first sampled window is spread across records/sections,
501            // - subsequent calls still rotate cyclically through the pool.
502            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    /// Returns true when `recipes` already contains the long-section auto recipe.
698    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    /// Stage A (source-level): decide whether to append auto recipe for `source`.
712    ///
713    /// Decision criteria:
714    /// - source must be eligible (`source_supports_chunk_pair_recipe`), and
715    /// - configured pool must not already include the auto recipe name.
716    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    /// Build the auto-injected long-section recipe.
726    ///
727    /// Semantics:
728    /// - Anchor selector: `Context`
729    /// - Positive selector: `Context`
730    /// - Negative selector: `Context` on a different record (`WrongArticle`)
731    ///
732    /// Anchor and positive are selected by two independent `select_chunk` calls
733    /// from the same context chunk candidate pool for the chosen record.
734    /// They are not concatenated and one is not derived from the other's text.
735    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    /// Resolve the source triplet plan.
749    ///
750    /// This is Stage A (source-level injection), not record-level execution.
751    ///
752    /// Algorithm for the auto-triplet injection path:
753    /// 1) Start from the configured source recipe pool (`configured_triplet_recipes_for_source`).
754    /// 2) Check auto-injection eligibility (`should_auto_inject_chunk_pair_recipe`), which is true only when:
755    ///    - the source has long sections discovered during ingest sync, and
756    ///    - the pool does not already include the auto recipe name.
757    /// 3) If eligible, append `source_chunk_pair_recipe`.
758    /// 4) Return both the effective recipe pool and whether step (3) happened.
759    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    /// Build a weighted, shuffled selection order from a slice of recipe weights.
790    ///
791    /// Each recipe with `weight > 0.0` receives a number of slots proportional to its weight
792    /// relative to the smallest positive weight, capped at
793    /// [`RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER`].  The slot list is then shuffled so that
794    /// deterministic round-robin cycling naturally draws recipes at the requested frequency:
795    ///
796    /// * `weight = [1.0, 1.0, 1.0]` → one slot each — identical to uniform round-robin.
797    /// * `weight = [3.0, 1.0]`       → three slots for recipe 0, one for recipe 1.
798    /// * `weight = [1.0, 0.0]`       → recipe 1 is excluded entirely (zero slots).
799    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    /// Weighted shuffle rotated by `rr_idx` for deterministic round-robin cycling.
808    ///
809    /// Calls [`Self::recipe_order_weighted_shuffled`] and rotates the result by
810    /// `rr_idx % order.len()` so successive batches start from different recipes,
811    /// preserving both the proportional frequency and the cross-batch cycling.
812    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    /// Weighted shuffled selection order for text recipes — same algorithm as
830    /// [`Self::recipe_order_weighted_shuffled`], kept separate so the text and triplet
831    /// round-robin counters can advance independently.
832    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        // The epoch advanced, changing the sampling permutation.
984        // emitted_text_hashes is NOT cleared: the record pool is unchanged,
985        // so texts already emitted are still off-limits.  The epoch
986        // changes the shuffling permutation, not the dedup window.
987        // Advance the epoch counter so source refreshes receive a different
988        // permutation seed (seed ^ epoch), producing a fresh traversal order.
989        self.epoch = self.epoch.saturating_add(1);
990        self.ingestion.set_epoch(self.epoch);
991        // Do NOT reset the epoch step here: it tracks the number of
992        // refresh batches, which is orthogonal to the sampling epoch.
993        // Resetting it would corrupt the persisted cursor value on resume.
994        // Reset per-source record cursors so the new epoch starts from the
995        // beginning of each source's permuted index space (not from where the
996        // previous epoch left off).
997        self.source_record_cursors.clear();
998        self.source_cycle_idx = 0;
999        // Reset the wrapped-tracking so all sources are eligible for sampling
1000        // in the new epoch.
1001        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                    // Fallback to any other record in the same split so split boundaries
1138                    // remain strictly isolated.
1139                    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                    // Fallback to any other record in the same split so split boundaries
1180                    // remain strictly isolated.
1181                    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                // Return the anchor record itself.  The caller
1211                // (finalize_triplet_with_negative) will call select_chunk on
1212                // this record using the recipe's negative_selector, which
1213                // picks from Context sections.  The dedup guard rejects
1214                // triplets where any two slots share identical text.
1215                Some((Arc::new(anchor_record.clone()), false))
1216            }
1217        }
1218    }
1219
1220    /// True when the recipe is the special auto-injected long-section chunk-pair recipe.
1221    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    /// Select anchor/positive chunk pair from selectors for a single record.
1226    ///
1227    /// When selectors are the same, retries until a valid pair is found or retry limit is hit.
1228    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    /// Select two distinct *window* chunks for the auto-injected recipe using recipe selectors.
1260    ///
1261    /// This method intentionally uses the recipe's selectors (no hardcoded selector values).
1262    /// For the current auto recipe those selectors are both `Role(Context)`, but the behavior
1263    /// remains tied to recipe configuration, not this helper body.
1264    ///
1265    /// NOTE: this does *not* decide whether the auto recipe is enabled for a source.
1266    /// Source-level enablement happens in `resolve_source_triplet_plan` via
1267    /// `should_auto_inject_chunk_pair_recipe`.
1268    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    /// Stage B (record-level): guard for the auto-injected recipe.
1280    ///
1281    /// Requires at least two `ChunkView::Window` candidates across sections addressed by
1282    /// `selector`. This makes the auto recipe eligibility check explicit at execution time.
1283    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, &section.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    /// Shared selector-pair stage for both standard and auto-injected triplet recipes.
1325    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        // Snapshot the raw anchor text for BM25 querying before decoration
1342        // added the metadata prefix — prefix tokens are absent from the index.
1343        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    /// Execute the special auto-injected long-section chunk-pair recipe.
1354    ///
1355    /// Algorithm:
1356    /// 0) Stage B check: current record must expose at least two window chunks for selector.
1357    /// 1) Draw anchor and positive from the same record using recipe selectors.
1358    /// 2) Require the two chunks to be distinct.
1359    /// 3) Additionally require both chunks to be `ChunkView::Window` variants.
1360    /// 4) Apply metadata decoration to both chunks.
1361    /// 5) Draw the negative according to recipe negative strategy.
1362    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    /// Finalize a triplet by selecting a negative and applying a deterministic 50 % coin-flip
1396    /// that swaps anchor and positive.
1397    ///
1398    /// # Anchor/positive swap
1399    ///
1400    /// Before constructing the [`SampleTriplet`], the sampler tests the least-significant bit of
1401    /// the next RNG word. When the bit is zero (≈ 50 % of the time) the anchor and positive
1402    /// slots are exchanged so that what was originally selected as the positive becomes the
1403    /// anchor, and vice-versa.
1404    ///
1405    /// **Why this matters:** contrastive objectives such as InfoNCE treat the two non-negative
1406    /// slots asymmetrically only in so far as downstream code does. If the model can learn a
1407    /// positional shortcut — for example "the anchor is always the shorter text" — it can
1408    /// achieve low loss without learning the intended similarity structure. The swap eliminates
1409    /// that opportunity by ensuring both orderings appear at equal frequency, forcing the model
1410    /// to treat the two positive-pair slots symmetrically.
1411    ///
1412    /// The negative chunk is unaffected by the swap.
1413    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        // 50 % coin-flip: swap anchor and positive to prevent positional shortcuts.
1432        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        // Reject if any two slots share identical rendered text.  This catches sources
1439        // that produce multiple records with the same string content: a negative whose
1440        // text matches the positive (or the anchor) would produce a trivially-invalid
1441        // training example.  String equality short-circuits on the first differing byte,
1442        // so the guard is effectively free when texts differ (the common case).
1443        //
1444        // When `allow_same_anchor_positive` is set the anchor==positive check is skipped
1445        // to support SimCSE-style training where both slots carry identical text and the
1446        // model's dropout layers produce the required embedding variation at train time.
1447        // The negative must still differ from both anchor and positive.
1448        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            // For the auto long-section recipe, use one coherence signal (proximity)
1519            // for the anchor/positive pair.
1520            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            // Keep negative weighting simple in this recipe: trust + floor only.
1526            return (pair_weight + pair_weight + negative_weight) / 3.0;
1527        }
1528        // Non-auto recipes also apply anchor-positive proximity, while keeping
1529        // negative weighting trust-only.
1530        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    // ── parallel-batch helpers (&self, rng-based, no cursor maps) ────────────
1577
1578    /// Select a chunk using rng instead of cursor maps (for parallel execution).
1579    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    /// Select a chunk by role using rng (parallel path — no role_cursors written).
1618    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    /// Decorate a chunk using a provided rng (parallel-safe, &self).
1648    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    /// Select an anchor/positive pair for the parallel path (rng-based).
1689    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        // No swap here — finalize_triplet_parallel applies the coin-flip swap.
1712        Some((anchor_chunk, positive_chunk))
1713    }
1714
1715    /// Select (anchor_chunk, positive_chunk, anchor_raw_text) for a recipe — parallel-safe.
1716    /// Does NOT select a negative; that happens sequentially in Phase 3.
1717    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, &section.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                // On first use for this (record,role), choose a deterministic hashed
1802                // section offset instead of always starting at the first matching section.
1803                // This avoids systematic head-bias while preserving reproducibility.
1804                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    /// Materialize chunk windows for one section according to `ChunkingStrategy`.
1820    ///
1821    /// Window layout algorithm:
1822    /// 1) Tokenize section text with `split_whitespace`.
1823    /// 2) Set `span = min(max_window_tokens, total_tokens)`.
1824    /// 3) If `span == total_tokens`, emit one full-section `ChunkView::Window`.
1825    /// 4) Otherwise, for each configured overlap:
1826    ///    - `stride = max(1, span - overlap)`
1827    ///    - emit windows `[start, min(start + span, total_tokens))` while advancing by `stride`.
1828    /// 5) Optionally append one `SummaryFallback` chunk when section is longer than `max_window_tokens`.
1829    ///
1830    /// This is the source of truth for how chunk windows are "laid" in a record.
1831    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    /// Stage A helper: true when record contains an Anchor/Context section whose token
1868    /// count exceeds `chunking.max_window_tokens`.
1869    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(&section.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        // Determine which records are being evicted BEFORE clearing self.records.
1885        // Use owned strings to avoid borrowing self.records across the clear().
1886        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        // Replace the record pool with the current cache snapshot so that
1891        // records evicted from the bounded in-memory cache no longer consume
1892        // memory here.  Without this clear, stale records accumulate across
1893        // cache refreshes, causing unbounded memory growth.
1894        self.records.clear();
1895        // Rebuild the long-section set from the new record pool.  Without
1896        // clearing first, sources whose records were evicted would leave
1897        // stale entries, though this is bounded by num_sources in practice.
1898        self.sources_with_long_sections.clear();
1899        // Prune emitted text hashes only for records that were actually
1900        // evicted from the pool, not on every ingestion advance.  This
1901        // preserves cross-batch dedup within a stable pool: with a window
1902        // of 512 and batch_size of 4, only 4 records are swapped per
1903        // advance, and the other 508 records' emitted hashes must survive.
1904        //
1905        // Old code cleared the entire set on ANY pool change, which broke
1906        // cross-batch dedup on every single advance call.
1907        for evicted_id in &evicted {
1908            self.emitted_text_hashes.remove(*evicted_id);
1909        }
1910        // Cursor state (BM25 hard-negative caches, chunk/role cursors) must
1911        // never outlive a record snapshot boundary.  BM25 backend clears its
1912        // negative cursors here; chunk/role cursors are pruned below in
1913        // prune_cursor_state() against the new record pool.
1914        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                // Mark source-level eligibility for auto-injected chunk-pair recipe.
1921                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        // Epoch tracking and source-state management must run every batch regardless
1963        // of whether the record pool changed — reconcile advances the sampling cursor.
1964        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    /// Ensure ingestion cursors (including epoch step) are loaded from the
1973    /// persisted state so the step counter reflects the saved value before
1974    /// the wrapper increments it for the current call.
1975    ///
1976    /// Only loads ingestion cursors and source epoch.  The sampling state
1977    /// (RNG, source_record_cursors, …) is restored later by
1978    /// `ingest_internal_for_split` → `ensure_source_state`, after
1979    /// `sync_records_from_cache` has populated `source_record_indices`.
1980    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    /// Helper that centralises the weight-fallback logic.
2041    ///
2042    /// If `weights` is `Some` and non-uniform or if `weights` is `None`,
2043    /// delegates to the appropriate ingest method.
2044    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    /// Select one triplet candidate for a specific source and split.
2098    ///
2099    /// This helper is the source-level recipe execution path used by both pair and
2100    /// triplet batch builders. It always starts from `triplet_recipes_for_source`,
2101    /// which already includes any auto-injected long-section recipe when eligible.
2102    ///
2103    /// Returns:
2104    /// - candidate triplet (with recipe) if one was sampled,
2105    /// - number of recipe attempts consumed (for round-robin bookkeeping).
2106    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        // Stage A (source-level injection): resolve effective recipe pool,
2115        // including auto-injected long-section recipe when source is eligible.
2116        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            // Stage B/C happen inside `make_triplet_with_anchor`:
2140            // - Stage B: record-level auto-recipe eligibility gate,
2141            // - Stage C: chunk window materialization/selection.
2142            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                        // Look up the anchor record's label to override structural defaults.
2215                        // For SRD pair-mode data, the label is stored in DataRecord.
2216                        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                    // Look up the anchor record's label to override structural defaults.
2311                    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    /// Check whether `sample` passes both per-batch and per-window text dedup.
2362    /// The per-batch set is never cleared mid-batch (survives epoch advances),
2363    /// while `emitted_text_hashes` is cleared on epoch advance to allow fresh
2364    /// permutations.  Both must pass for the sample to be emitted.
2365    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        // ── Declarations for multi-source parallel path ────────────────────────
2601        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        // ── Phase 1: sequential anchor pre-assignment ──────────────────────────
2617        // Advance cursors and pick anchors deterministically. Capture a fork seed
2618        // for each slot so that Phase 2 can reconstruct independent rngs in
2619        // deterministic order.
2620        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        // ── Phase 2: parallel anchor+positive selection ────────────────────────
2656        // Only chunk selection runs in parallel; negative selection and the
2657        // coin-flip swap remain sequential in Phase 3 so that the BM25 cursor
2658        // rotation (shared RwLock state) is deterministic across runs.
2659        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                // Use weighted shuffle so recipes with higher weight are tried more
2672                // often and zero-weight recipes are excluded entirely.
2673                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        // ── Phase 3: sequential finalization (negative selection + coin-flip) ───
2696        // Sort by slot_idx to restore deterministic ordering, then finalize each
2697        // candidate sequentially so that BM25 cursor rotation advances in a fixed
2698        // order regardless of rayon's thread scheduling.
2699        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            // 50 % coin-flip swap using the pre-sampled per-slot seed + main rng.
2726            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    // ── test-only helpers ─────────────────────────────────────────────────────
2785
2786    /// Return a mutable reference to the concrete [`backends::Bm25Backend`], panicking
2787    /// when the feature is disabled (which won't happen — call sites are
2788    /// equally gated on `#[cfg(feature = "bm25-mining")]`).
2789    #[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    /// Test shim: `recipe_order_weighted_shuffled` using the inner rng.
2799    #[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    /// Test shim: `recipe_order_weighted_cycled` using the inner rng.
2808    #[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    /// Test shim: `text_recipe_order_weighted_shuffled` using the inner rng.
2821    #[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    /// Test shim: `text_recipe_order_weighted_cycled` using the inner rng.
2830    #[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    /// Test shim: `select_negative_record` using the inner rng.
2843    #[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    /// Test shim: `make_triplet_with_anchor` using the inner rng.
2858    #[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    /// Test shim: `decorate_chunk` using the inner rng.
2871    #[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    // ── extended-metrics helpers ───────────────────────────────────────────────
2879
2880    /// Return cumulative `(fallback_count, selection_count)` from the BM25
2881    /// backend since it was created.  Only available when both `bm25-mining`
2882    /// and `extended-metrics` features are active.
2883    #[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    /// Look up `anchor`'s split label and return BM25-ranked candidate IDs.
2889    ///
2890    /// Mirrors the old `bm25_ranked_candidates` method that tests called on
2891    /// `TripletSamplerInner` before BM25 state moved into `Bm25Backend`.
2892    #[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
2907/// Build a shuffled selection order from a slice of recipe weights.
2908///
2909/// Each recipe with `weight > 0.0` receives a number of slots proportional to its weight
2910/// relative to the smallest positive weight, capped at
2911/// [`RECIPE_ORDER_MAX_WEIGHT_MULTIPLIER`].  The resulting expanded slot list is then
2912/// shuffled, giving each recipe approximately the requested frequency in round-robin
2913/// cycling.
2914///
2915/// This is a free function (no `self`) so it can be called from rayon parallel closures as
2916/// well as from within the `TripletSamplerInner` instance methods.
2917fn 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    /// Create a sampler from config and a split-state backend.
2961    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    /// Create a sampler from config with a custom chunking implementation.
2969    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    /// Return an unweighted pair batch for `split`.
2981    ///
2982    /// Samples all sources uniformly. Prefer
2983    /// [`Self::next_pair_batch_with_weights_for_split`] with an explicit
2984    /// per-source weight map to honor a data mixture.
2985    #[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    /// Return an unweighted text batch for `split`.
2996    ///
2997    /// Samples all sources uniformly. Prefer
2998    /// [`Self::next_text_batch_with_weights_for_split`] with an explicit
2999    /// per-source weight map to honor a data mixture.
3000    #[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    /// Return an unweighted triplet batch for `split`.
3008    ///
3009    /// Samples all sources uniformly. Prefer
3010    /// [`Self::next_triplet_batch_with_weights_for_split`] with an explicit
3011    /// per-source weight map to honor a data mixture.
3012    #[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    /// Shared retry loop for all three batch types: load state, bump step,
3023    /// retry on exhaustion up to EXHAUSTION_RETRY_LIMIT times, and return
3024    /// the batch or propagate the final error.
3025    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        // Load persisted state BEFORE bumping step so the restored step
3038        // counter is reflected before we increment it for this call.
3039        inner.ensure_ingestion_cursors_loaded()?;
3040        // Bump epoch step once per public call, regardless of success/exhaustion.
3041        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        // The for-loop always returns via the match arms above.
3052        unreachable!()
3053    }
3054
3055    /// Return a weighted pair batch for `split` using per-source weights.
3056    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    /// Return a weighted text batch for `split` using per-source weights.
3067    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    /// Return a weighted triplet batch for `split` using per-source weights.
3078    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    /// Spawn a background prefetcher for triplet batches.
3089    ///
3090    /// Samples all sources uniformly. Prefer
3091    /// [`Self::prefetch_triplet_batches_with_weights`] with an explicit
3092    /// per-source weight map to honor a data mixture.
3093    #[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    /// Spawn a background prefetcher for weighted triplet batches.
3107    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    /// Spawn a background prefetcher for pair batches.
3119    ///
3120    /// Samples all sources uniformly. Prefer
3121    /// [`Self::prefetch_pair_batches_with_weights`] with an explicit
3122    /// per-source weight map to honor a data mixture.
3123    #[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    /// Spawn a background prefetcher for weighted pair batches.
3137    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    /// Spawn a background prefetcher for text batches.
3149    ///
3150    /// Samples all sources uniformly. Prefer
3151    /// [`Self::prefetch_text_batches_with_weights`] with an explicit
3152    /// per-source weight map to honor a data mixture.
3153    #[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    /// Spawn a background prefetcher for weighted text batches.
3167    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    /// Return the currently active text recipes.
3179    pub fn text_recipes(&self) -> Vec<TextRecipe> {
3180        let inner = self.inner.lock().unwrap();
3181        inner.text_recipes().to_vec()
3182    }
3183
3184    /// Register a data source for ingestion and sampling.
3185    ///
3186    /// Returns an error if the source's `id()` matches the reserved `__*__`
3187    /// pattern used for internal synthetic/metadata source identifiers.
3188    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    /// Force sampler epoch to `epoch` (advanced deterministic replay control).
3197    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    /// Persist sampler and split runtime state for restart-resume.
3203    ///
3204    /// When `save_to` is `Some(path)`, current persisted runtime state is also
3205    /// mirrored to `path` when supported by the split-store backend.
3206    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    /// Return cumulative `(fallback_count, selection_count)` from the BM25
3212    /// backend since it was created.  Only available when both `bm25-mining`
3213    /// and `extended-metrics` features are active.
3214    ///
3215    /// `selection_count` counts calls where the pool was non-empty.
3216    /// `fallback_count` counts the subset where BM25 produced no candidates
3217    /// intersecting the pool and random selection was used instead.
3218    #[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/// Dedup key for text batch generation: (record_id, text).
3296///
3297/// Unlike `triplet_chunk_key` which compares structural identity
3298/// (record_id, section_idx, window_index), this key compares actual
3299/// text content. This prevents text-columns mode sources (where Anchor
3300/// and Context sections carry identical text) from producing duplicate
3301/// samples from the same record.
3302#[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;