rs-chunks 0.1.0

Fast, high-fidelity document chunking for RAG — a pure-Rust engine covering 36 file formats (Office, OpenDocument, PDF, email, ebooks, notebooks, and more).
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

//! Legacy PowerPoint binary (`.ppt`, OLE/CFB) chunking. Reuses the `.doc`
//! paragraph chunkers and metadata on paragraphs extracted from the PowerPoint
//! Document stream.

pub mod cfb_reader;
pub mod images;
pub mod records;
pub mod structural;
pub mod text_extractor;
pub mod to_markdown;

use crate::chunk::Chunk;
use crate::error::{ChunkError, Result};
use crate::formats::doc::records_to_chunks;
use crate::formats::doc::structural::{
    build_page_aware_chunks, build_section_chunks, build_semantic_chunks, build_sentence_chunks,
    build_sliding_window_chunks, build_structural_chunks,
};
use crate::options::{ChunkMode, ChunkOptions};
use structural::{load_ppt_paragraphs, validate_ppt_path};

pub fn chunk(
    file_path: &str,
    mode: &str,
    window_size: usize,
    overlap: usize,
    sentences_per_chunk: usize,
    paragraphs_per_page: usize,
) -> Result<Vec<Chunk>> {
    validate_ppt_path(file_path).map_err(ChunkError::InvalidArg)?;
    let paragraphs = load_ppt_paragraphs(file_path).map_err(ChunkError::Parse)?;
    Ok(records_to_chunks(file_path, build_by_mode(paragraphs, mode, window_size, overlap, sentences_per_chunk, paragraphs_per_page)?))
}

fn build_by_mode(
    paragraphs: Vec<crate::formats::doc::text_extractor::DocParagraph>,
    mode: &str,
    window_size: usize,
    overlap: usize,
    sentences_per_chunk: usize,
    paragraphs_per_page: usize,
) -> Result<Vec<crate::formats::doc::structural::ChunkRecord>> {
    Ok(match mode {
        "default" | "structural" => build_structural_chunks(paragraphs),
        "section" => build_section_chunks(paragraphs),
        "semantic" => build_semantic_chunks(paragraphs),
        "sentence" => build_sentence_chunks(paragraphs, sentences_per_chunk),
        "page_aware" => build_page_aware_chunks(paragraphs, paragraphs_per_page),
        "sliding_window" => build_sliding_window_chunks(paragraphs, window_size, overlap),
        other => return Err(ChunkError::InvalidArg(format!("Unknown PPT mode: {other}"))),
    })
}

/// No-filesystem entry (wasm/browser).
pub fn chunk_from_bytes(data: &[u8], source: &str, mode: &str, window_size: usize, overlap: usize, sentences_per_chunk: usize, paragraphs_per_page: usize) -> Result<Vec<Chunk>> {
    let paragraphs = structural::load_ppt_paragraphs_bytes(data).map_err(ChunkError::Parse)?;
    Ok(records_to_chunks(source, build_by_mode(paragraphs, mode, window_size, overlap, sentences_per_chunk, paragraphs_per_page)?))
}

pub fn to_markdown_from_bytes(data: &[u8]) -> Result<String> {
    let paragraphs = structural::load_ppt_paragraphs_bytes(data).map_err(ChunkError::Parse)?;
    Ok(crate::formats::doc::to_markdown::render_paragraphs_markdown(paragraphs))
}

pub fn chunk_with_options(file_path: &str, opts: &ChunkOptions) -> Result<Vec<Chunk>> {
    let mode = match opts.mode {
        ChunkMode::Default => "default",
        ChunkMode::Structural => "structural",
        ChunkMode::Section => "section",
        ChunkMode::Semantic => "semantic",
        ChunkMode::Sentence => "sentence",
        ChunkMode::PageAware => "page_aware",
        ChunkMode::SlidingWindow => "sliding_window",
        other => {
            return Err(ChunkError::InvalidArg(format!(
                "PPT does not support mode '{}'",
                other.as_str()
            )))
        }
    };
    chunk(
        file_path,
        mode,
        opts.window_size,
        opts.overlap,
        opts.sentences_per_chunk,
        opts.paragraphs_per_page,
    )
}

/// Chunk with embedded images (ODRAW/blip extraction): image chunks first, then
/// text chunks, with indices renumbered across the combined list.
pub fn chunk_with_images(
    file_path: &str,
    mode: &str,
    window_size: usize,
    overlap: usize,
    sentences_per_chunk: usize,
    paragraphs_per_page: usize,
) -> Result<(Vec<Chunk>, Vec<(String, Vec<u8>)>)> {
    validate_ppt_path(file_path).map_err(ChunkError::InvalidArg)?;
    let res = match mode {
        "default" | "structural" => images::chunk_with_images_impl(file_path, build_structural_chunks),
        "section" => images::chunk_with_images_impl(file_path, build_section_chunks),
        "semantic" => images::chunk_with_images_impl(file_path, build_semantic_chunks),
        "sentence" => images::chunk_with_images_impl(file_path, |p| build_sentence_chunks(p, sentences_per_chunk)),
        "page_aware" => images::chunk_with_images_impl(file_path, |p| build_page_aware_chunks(p, paragraphs_per_page)),
        "sliding_window" => images::chunk_with_images_impl(file_path, |p| build_sliding_window_chunks(p, window_size, overlap)),
        other => return Err(ChunkError::InvalidArg(format!("Unknown PPT mode: {other}"))),
    };
    res.map_err(ChunkError::Parse)
}

pub fn to_markdown(file_path: &str) -> Result<String> {
    validate_ppt_path(file_path).map_err(ChunkError::InvalidArg)?;
    let paragraphs = load_ppt_paragraphs(file_path).map_err(ChunkError::Parse)?;
    Ok(crate::formats::doc::to_markdown::render_paragraphs_markdown(paragraphs))
}

pub fn to_markdown_with_images(file_path: &str) -> Result<(String, Vec<(String, Vec<u8>)>)> {
    to_markdown::to_markdown_with_images(file_path).map_err(ChunkError::Parse)
}

/// No-filesystem `chunk_with_images` (wasm/browser). `filename` is recorded as
/// each chunk's `source`.
pub fn chunk_with_images_from_bytes(bytes: &[u8], filename: &str, mode: &str, window_size: usize, overlap: usize, sentences_per_chunk: usize, paragraphs_per_page: usize) -> Result<(Vec<Chunk>, Vec<(String, Vec<u8>)>)> {
    let res = match mode {
        "default" | "structural" => images::chunk_with_images_impl_bytes(bytes, filename, build_structural_chunks),
        "section" => images::chunk_with_images_impl_bytes(bytes, filename, build_section_chunks),
        "semantic" => images::chunk_with_images_impl_bytes(bytes, filename, build_semantic_chunks),
        "sentence" => images::chunk_with_images_impl_bytes(bytes, filename, |p| build_sentence_chunks(p, sentences_per_chunk)),
        "page_aware" => images::chunk_with_images_impl_bytes(bytes, filename, |p| build_page_aware_chunks(p, paragraphs_per_page)),
        "sliding_window" => images::chunk_with_images_impl_bytes(bytes, filename, |p| build_sliding_window_chunks(p, window_size, overlap)),
        other => return Err(ChunkError::InvalidArg(format!("Unknown PPT mode: {other}"))),
    };
    res.map_err(ChunkError::Parse)
}

pub fn to_markdown_with_images_from_bytes(bytes: &[u8]) -> Result<(String, Vec<(String, Vec<u8>)>)> {
    to_markdown::to_markdown_with_images_bytes(bytes).map_err(ChunkError::Parse)
}

pub fn stream(
    file_path: &str,
    mode: &str,
    window_size: usize,
    overlap: usize,
    sentences_per_chunk: usize,
    paragraphs_per_page: usize,
) -> Result<impl Iterator<Item = Result<Chunk>>> {
    Ok(chunk(file_path, mode, window_size, overlap, sentences_per_chunk, paragraphs_per_page)?
        .into_iter()
        .map(Ok))
}