bbpe/
corpus.rs

1//! Facilities for discovering input files and loading binary corpora.
2
3use std::fs::File;
4use std::io::{BufRead, BufReader, Read};
5use std::path::{Path, PathBuf};
6
7use flate2::read::MultiGzDecoder;
8use serde_json::Value;
9use walkdir::WalkDir;
10
11use crate::config::IngestConfig;
12use crate::error::{BbpeError, Result};
13
14/// Specification describing a JSONL input file and the nested field to extract per record.
15#[derive(Debug, Clone)]
16pub struct JsonlSpec {
17    /// Path to the JSONL file.
18    pub path: PathBuf,
19    /// Nested field path (e.g. `["choices", "0", "text"]`).
20    pub field_path: Vec<String>,
21}
22
23impl std::str::FromStr for JsonlSpec {
24    type Err = String;
25
26    fn from_str(spec: &str) -> std::result::Result<Self, Self::Err> {
27        let mut parts = spec.rsplitn(2, ':');
28        let field = parts
29            .next()
30            .ok_or_else(|| "jsonl specification must include `PATH:FIELD`".to_string())?;
31        let path = parts
32            .next()
33            .ok_or_else(|| "jsonl specification missing file path".to_string())?;
34        if path.is_empty() {
35            return Err("jsonl file path cannot be empty".into());
36        }
37        if field.is_empty() {
38            return Err("jsonl field path cannot be empty".into());
39        }
40        let field_path: Vec<String> = field
41            .split('.')
42            .map(str::trim)
43            .filter(|segment| !segment.is_empty())
44            .map(|segment| segment.to_string())
45            .collect();
46        if field_path.is_empty() {
47            return Err("jsonl field path must contain at least one segment".into());
48        }
49        Ok(Self {
50            path: PathBuf::from(path),
51            field_path,
52        })
53    }
54}
55
56/// Lazily streams binary corpus chunks from disk without retaining them all in memory.
57pub struct BinaryChunkStream {
58    files: Vec<PathBuf>,
59    cfg: IngestConfig,
60    file_index: usize,
61    current_file: Option<File>,
62    current_path: Option<PathBuf>,
63    total_chunks: usize,
64    total_bytes: usize,
65}
66
67/// Discovers files rooted at the provided input paths according to the ingest configuration.
68///
69/// Directories are traversed recursively by default; set [`IngestConfig::recursive`] to `false`
70/// to limit discovery to the first level.  Symlink traversal is controlled through
71/// [`IngestConfig::follow_symlinks`].
72pub fn collect_paths<P: AsRef<Path>>(inputs: &[P], cfg: &IngestConfig) -> Result<Vec<PathBuf>> {
73    let mut files = Vec::new();
74    for input in inputs {
75        let path = input.as_ref();
76        if !path.exists() {
77            return Err(BbpeError::InvalidConfig(format!(
78                "input path {path:?} does not exist"
79            )));
80        }
81        let metadata = path
82            .symlink_metadata()
83            .map_err(|err| BbpeError::io(err, Some(path.to_path_buf())))?;
84        if metadata.is_dir() {
85            if cfg.recursive {
86                let walker = WalkDir::new(path).follow_links(cfg.follow_symlinks);
87                for entry in walker {
88                    let entry = entry.map_err(|err| BbpeError::Internal(err.to_string()))?;
89                    if entry.file_type().is_file() {
90                        files.push(entry.path().to_path_buf());
91                    }
92                }
93            } else {
94                for entry in std::fs::read_dir(path)
95                    .map_err(|err| BbpeError::io(err, Some(path.to_path_buf())))?
96                {
97                    let entry =
98                        entry.map_err(|err| BbpeError::io(err, Some(path.to_path_buf())))?;
99                    let entry_path = entry.path();
100                    if entry_path.is_file() {
101                        files.push(entry_path);
102                    }
103                }
104            }
105        } else if metadata.is_file() {
106            files.push(path.to_path_buf());
107        }
108    }
109    if files.is_empty() {
110        return Err(BbpeError::InvalidConfig(
111            "no files discovered in provided inputs".into(),
112        ));
113    }
114    Ok(files)
115}
116
117/// Loads binary corpora into memory chunks based on the ingest configuration.
118///
119/// Files are loaded in-order, optionally split into fixed-size chunks configured via
120/// [`IngestConfig::chunk_size`].  Empty chunks are discarded to avoid degenerate training input.
121pub fn load_binary_corpus<P: AsRef<Path>>(
122    inputs: &[P],
123    cfg: &IngestConfig,
124) -> Result<Vec<Vec<u8>>> {
125    let file_paths = collect_paths(inputs, cfg)?;
126    let mut sequences = Vec::new();
127    for file_path in file_paths {
128        let mut file =
129            File::open(&file_path).map_err(|err| BbpeError::io(err, Some(file_path.clone())))?;
130        if cfg.chunk_size == 0 {
131            let mut buffer = Vec::new();
132            file.read_to_end(&mut buffer)
133                .map_err(|err| BbpeError::io(err, Some(file_path.clone())))?;
134            if !buffer.is_empty() {
135                sequences.push(buffer);
136            }
137            continue;
138        }
139
140        loop {
141            let mut buffer = vec![0u8; cfg.chunk_size];
142            let read = file
143                .read(&mut buffer)
144                .map_err(|err| BbpeError::io(err, Some(file_path.clone())))?;
145            if read == 0 {
146                break;
147            }
148            buffer.truncate(read);
149            sequences.push(buffer);
150        }
151    }
152    if sequences.is_empty() {
153        return Err(BbpeError::InvalidConfig(
154            "no binary data could be loaded from inputs".into(),
155        ));
156    }
157    Ok(sequences)
158}
159
160/// Loads newline-delimited JSON (JSONL) records and extracts byte sequences from a specified field.
161pub fn load_jsonl_corpus(specs: &[JsonlSpec]) -> Result<Vec<Vec<u8>>> {
162    if specs.is_empty() {
163        return Ok(Vec::new());
164    }
165
166    let sequences: Result<Vec<Vec<u8>>> = stream_jsonl_corpus(specs)?.collect();
167    let sequences = sequences?;
168    if sequences.is_empty() {
169        return Err(BbpeError::InvalidConfig(
170            "no sequences extracted from JSONL inputs".into(),
171        ));
172    }
173    Ok(sequences)
174}
175
176fn open_jsonl_reader(path: &Path) -> Result<Box<dyn BufRead + Send>> {
177    let file = File::open(path).map_err(|err| BbpeError::io(err, Some(path.to_path_buf())))?;
178    if is_gzip_path(path) {
179        let decoder = MultiGzDecoder::new(file);
180        Ok(Box::new(BufReader::new(decoder)))
181    } else {
182        Ok(Box::new(BufReader::new(file)))
183    }
184}
185
186fn is_gzip_path(path: &Path) -> bool {
187    path.extension()
188        .and_then(|ext| ext.to_str())
189        .map(|ext| ext.eq_ignore_ascii_case("gz"))
190        .unwrap_or(false)
191}
192
193/// Streams binary corpus chunks from disk, avoiding materialising the entire corpus at once.
194///
195/// The returned iterator yields `Vec<u8>` buffers according to [`IngestConfig::chunk_size`], skipping
196/// empty results.  File discovery honours the same semantics as [`load_binary_corpus`].  Metadata about
197/// the total number of chunks and bytes is precomputed from filesystem information so callers can set
198/// up progress reporting without forcing eager loading.
199pub fn stream_binary_corpus<P: AsRef<Path>>(
200    inputs: &[P],
201    cfg: &IngestConfig,
202) -> Result<BinaryChunkStream> {
203    let files = collect_paths(inputs, cfg)?;
204    let mut total_chunks = 0usize;
205    let mut total_bytes = 0usize;
206    let chunk_size = cfg.chunk_size;
207
208    for path in &files {
209        let metadata = path
210            .metadata()
211            .map_err(|err| BbpeError::io(err, Some(path.clone())))?;
212        let file_len = metadata.len();
213        let file_len_usize = usize::try_from(file_len).map_err(|_| {
214            BbpeError::InvalidConfig(format!(
215                "input file {} exceeds usize::MAX ({} bytes)",
216                path.display(),
217                file_len
218            ))
219        })?;
220        total_bytes = total_bytes.saturating_add(file_len_usize);
221
222        if file_len == 0 {
223            continue;
224        }
225
226        if chunk_size == 0 {
227            total_chunks = total_chunks.saturating_add(1);
228        } else {
229            let size = u64::try_from(chunk_size).map_err(|_| {
230                BbpeError::InvalidConfig(format!("chunk size {chunk_size} exceeds u64::MAX"))
231            })?;
232            let mut chunks = file_len / size;
233            if file_len % size != 0 {
234                chunks = chunks.saturating_add(1);
235            }
236            let chunks_usize = usize::try_from(chunks).map_err(|_| {
237                BbpeError::InvalidConfig(format!(
238                    "file {} produces more chunks than usize::MAX allows",
239                    path.display()
240                ))
241            })?;
242            total_chunks = total_chunks.saturating_add(chunks_usize);
243        }
244    }
245
246    if total_chunks == 0 {
247        return Err(BbpeError::InvalidConfig(
248            "no binary data could be loaded from inputs".into(),
249        ));
250    }
251
252    Ok(BinaryChunkStream {
253        files,
254        cfg: cfg.clone(),
255        file_index: 0,
256        current_file: None,
257        current_path: None,
258        total_chunks,
259        total_bytes,
260    })
261}
262
263impl BinaryChunkStream {
264    /// Returns the total number of chunks that will be produced.
265    pub fn total_chunks(&self) -> usize {
266        self.total_chunks
267    }
268
269    /// Returns the aggregate byte count across all inputs.
270    pub fn total_bytes(&self) -> usize {
271        self.total_bytes
272    }
273}
274
275impl Iterator for BinaryChunkStream {
276    type Item = Result<Vec<u8>>;
277
278    fn next(&mut self) -> Option<Self::Item> {
279        loop {
280            if self.file_index >= self.files.len() {
281                return None;
282            }
283
284            if self.current_file.is_none() {
285                let path = self.files[self.file_index].clone();
286                match File::open(&path) {
287                    Ok(file) => {
288                        self.current_path = Some(path);
289                        self.current_file = Some(file);
290                    }
291                    Err(err) => {
292                        self.file_index = self.file_index.saturating_add(1);
293                        return Some(Err(BbpeError::io(err, Some(path))));
294                    }
295                }
296            }
297
298            let chunk_size = self.cfg.chunk_size;
299            if chunk_size == 0 {
300                let mut file = self
301                    .current_file
302                    .take()
303                    .expect("current file must exist when chunk_size == 0");
304                let path = self
305                    .current_path
306                    .take()
307                    .unwrap_or_else(|| self.files[self.file_index].clone());
308                let mut buffer = Vec::new();
309                let result = file.read_to_end(&mut buffer);
310                self.file_index = self.file_index.saturating_add(1);
311                match result {
312                    Ok(_) => {
313                        if buffer.is_empty() {
314                            continue;
315                        }
316                        return Some(Ok(buffer));
317                    }
318                    Err(err) => return Some(Err(BbpeError::io(err, Some(path)))),
319                }
320            } else {
321                let file = self
322                    .current_file
323                    .as_mut()
324                    .expect("current file must exist when chunk_size > 0");
325                let path = self
326                    .current_path
327                    .as_ref()
328                    .cloned()
329                    .unwrap_or_else(|| self.files[self.file_index].clone());
330
331                let mut buffer = vec![0u8; chunk_size];
332                match file.read(&mut buffer) {
333                    Ok(0) => {
334                        self.current_file = None;
335                        self.current_path = None;
336                        self.file_index = self.file_index.saturating_add(1);
337                        continue;
338                    }
339                    Ok(read) => {
340                        buffer.truncate(read);
341                        if buffer.is_empty() {
342                            continue;
343                        }
344                        return Some(Ok(buffer));
345                    }
346                    Err(err) => {
347                        self.current_file = None;
348                        self.current_path = None;
349                        self.file_index = self.file_index.saturating_add(1);
350                        return Some(Err(BbpeError::io(err, Some(path))));
351                    }
352                }
353            }
354        }
355    }
356}
357
358/// Streams records from JSONL files, yielding extracted byte sequences without buffering everything.
359pub fn stream_jsonl_corpus(specs: &[JsonlSpec]) -> Result<JsonlStream> {
360    Ok(JsonlStream::new(specs.to_vec()))
361}
362
363/// Iterator over JSONL records that extracts the configured nested field per line.
364pub struct JsonlStream {
365    specs: Vec<JsonlSpec>,
366    current_index: usize,
367    reader: Option<Box<dyn BufRead + Send>>,
368    buffer: String,
369    line_index: usize,
370}
371
372impl JsonlStream {
373    fn new(specs: Vec<JsonlSpec>) -> Self {
374        Self {
375            specs,
376            current_index: 0,
377            reader: None,
378            buffer: String::new(),
379            line_index: 0,
380        }
381    }
382}
383
384impl Iterator for JsonlStream {
385    type Item = Result<Vec<u8>>;
386
387    fn next(&mut self) -> Option<Self::Item> {
388        loop {
389            if self.current_index >= self.specs.len() {
390                return None;
391            }
392
393            if self.reader.is_none() {
394                let spec = &self.specs[self.current_index];
395                match open_jsonl_reader(&spec.path) {
396                    Ok(reader) => {
397                        self.reader = Some(reader);
398                        self.line_index = 0;
399                    }
400                    Err(err) => {
401                        self.current_index = self.current_index.saturating_add(1);
402                        return Some(Err(err));
403                    }
404                }
405            }
406
407            let spec = &self.specs[self.current_index];
408            let reader = self
409                .reader
410                .as_mut()
411                .expect("reader should be initialised before reading");
412            self.buffer.clear();
413            match reader.read_line(&mut self.buffer) {
414                Ok(0) => {
415                    self.reader = None;
416                    self.line_index = 0;
417                    self.current_index = self.current_index.saturating_add(1);
418                    continue;
419                }
420                Ok(_) => {
421                    self.line_index = self.line_index.saturating_add(1);
422                    if self.buffer.trim().is_empty() {
423                        continue;
424                    }
425                    match parse_jsonl_line(&self.buffer, spec, self.line_index) {
426                        Ok(Some(bytes)) => return Some(Ok(bytes)),
427                        Ok(None) => continue,
428                        Err(err) => {
429                            self.reader = None;
430                            self.current_index = self.current_index.saturating_add(1);
431                            return Some(Err(err));
432                        }
433                    }
434                }
435                Err(err) => {
436                    let path = spec.path.clone();
437                    self.reader = None;
438                    self.current_index = self.current_index.saturating_add(1);
439                    return Some(Err(BbpeError::io(err, Some(path))));
440                }
441            }
442        }
443    }
444}
445
446fn parse_jsonl_line(line: &str, spec: &JsonlSpec, line_idx: usize) -> Result<Option<Vec<u8>>> {
447    let value: Value = serde_json::from_str(line).map_err(|err| {
448        BbpeError::InvalidConfig(format!(
449            "failed to parse JSON in {} on line {}: {err}",
450            spec.path.display(),
451            line_idx
452        ))
453    })?;
454    let mut current = &value;
455    for key in &spec.field_path {
456        current = current.get(key).ok_or_else(|| {
457            BbpeError::InvalidConfig(format!(
458                "field `{}` missing in {} on line {}",
459                spec.field_path.join("."),
460                spec.path.display(),
461                line_idx
462            ))
463        })?;
464    }
465    let text = current.as_str().ok_or_else(|| {
466        BbpeError::InvalidConfig(format!(
467            "field `{}` in {} line {} is not a string",
468            spec.field_path.join("."),
469            spec.path.display(),
470            line_idx
471        ))
472    })?;
473    if text.is_empty() {
474        return Ok(None);
475    }
476    Ok(Some(text.as_bytes().to_vec()))
477}
478
479#[cfg(test)]
480mod tests {
481    use super::*;
482    use std::fs;
483    use tempfile::tempdir;
484
485    #[test]
486    fn collect_paths_discovers_files_recursively() {
487        let dir = tempdir().expect("tempdir");
488        let nested = dir.path().join("nested");
489        fs::create_dir(&nested).expect("create nested directory");
490        let file_a = dir.path().join("a.bin");
491        let file_b = nested.join("b.bin");
492        fs::write(&file_a, [1u8, 2, 3]).expect("write a");
493        fs::write(&file_b, [4u8, 5, 6]).expect("write b");
494
495        let cfg = IngestConfig {
496            recursive: true,
497            ..IngestConfig::default()
498        };
499        let mut paths = collect_paths(&[dir.path()], &cfg).expect("collect paths");
500        paths.sort();
501        assert_eq!(paths, vec![file_a, file_b]);
502    }
503
504    #[test]
505    fn load_binary_corpus_splits_chunks() {
506        let dir = tempdir().expect("tempdir");
507        let file = dir.path().join("data.bin");
508        let bytes: Vec<u8> = (0..=9).collect();
509        fs::write(&file, &bytes).expect("write data");
510
511        let cfg = IngestConfig {
512            chunk_size: 4,
513            ..IngestConfig::default()
514        };
515        let sequences =
516            load_binary_corpus(&[file], &cfg).expect("load corpus with chunking enabled");
517        assert_eq!(sequences.len(), 3);
518        assert_eq!(sequences[0], vec![0, 1, 2, 3]);
519        assert_eq!(sequences[1], vec![4, 5, 6, 7]);
520        assert_eq!(sequences[2], vec![8, 9]);
521    }
522
523    #[test]
524    fn load_binary_corpus_entire_file_when_chunk_zero() {
525        let dir = tempdir().expect("tempdir");
526        let file = dir.path().join("data.bin");
527        let bytes: Vec<u8> = (0..=9).collect();
528        fs::write(&file, &bytes).expect("write data");
529
530        let cfg = IngestConfig {
531            chunk_size: 0,
532            ..IngestConfig::default()
533        };
534        let sequences = load_binary_corpus(&[file], &cfg).expect("load corpus without chunking");
535        assert_eq!(sequences, vec![bytes]);
536    }
537
538    #[test]
539    fn stream_binary_corpus_matches_loaded_sequences() {
540        let dir = tempdir().expect("tempdir");
541        let file = dir.path().join("data.bin");
542        let bytes: Vec<u8> = (0..=15).collect();
543        fs::write(&file, &bytes).expect("write data");
544
545        let cfg = IngestConfig {
546            chunk_size: 4,
547            ..IngestConfig::default()
548        };
549        let expected = load_binary_corpus(std::slice::from_ref(&file), &cfg)
550            .expect("load corpus with chunking enabled");
551        let stream =
552            stream_binary_corpus(&[file], &cfg).expect("stream corpus with chunking enabled");
553        assert_eq!(stream.total_chunks(), expected.len());
554        assert_eq!(
555            stream.total_bytes(),
556            expected.iter().map(|chunk| chunk.len()).sum::<usize>()
557        );
558        let streamed = stream
559            .map(|item| item.expect("stream item"))
560            .collect::<Vec<_>>();
561        assert_eq!(streamed, expected);
562    }
563
564    #[test]
565    fn stream_jsonl_corpus_matches_loader() {
566        let dir = tempdir().expect("tempdir");
567        let file = dir.path().join("data.jsonl");
568        let contents = r#"
569{"text":"alpha"}
570{"text":""}
571{"text":"beta gamma"}
572"#;
573        fs::write(&file, contents.trim()).expect("write jsonl");
574        let spec = JsonlSpec {
575            path: file.clone(),
576            field_path: vec!["text".to_string()],
577        };
578
579        let expected =
580            load_jsonl_corpus(std::slice::from_ref(&spec)).expect("load jsonl corpus eagerly");
581        assert_eq!(expected.len(), 2);
582
583        let streamed: Result<Vec<_>> = stream_jsonl_corpus(std::slice::from_ref(&spec))
584            .expect("build stream")
585            .collect();
586        assert_eq!(streamed.expect("collect stream"), expected);
587    }
588}
bbpe/corpus.rs

bbpe/
corpus.rs
