kiru 0.1.11

Fast text chunking for Rust
Documentation 
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

// kiru-core/src/bin/benchmark.rs

use kiru::{ChunkerBuilder, Source};
use serde::Serialize;
use std::env;
use std::time::Instant;

#[derive(Serialize)]
struct BenchmarkResult {
    elapsed_secs: f64,
    num_chunks: usize,
    total_bytes: usize,
    throughput_mb_s: f64,
}

#[derive(Serialize)]
struct BenchmarkError {
    error: String,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    if args.len() < 6 {
        let error = BenchmarkError {
            error: "Usage: benchmark <strategy> <source_type> <path> <chunk_size> <overlap>"
                .to_string(),
        };
        eprintln!("{}", serde_json::to_string(&error).unwrap());
        std::process::exit(1);
    }

    let strategy = &args[1]; // "bytes" or "chars"
    let source_type = &args[2]; // "string" or "file" or "http" or "glob"
    let path = &args[3];
    let chunk_size: usize = match args[4].parse() {
        Ok(v) => v,
        Err(e) => {
            let error = BenchmarkError {
                error: format!("Invalid chunk_size: {}", e),
            };
            eprintln!("{}", serde_json::to_string(&error).unwrap());
            std::process::exit(1);
        }
    };
    let overlap: usize = match args[5].parse() {
        Ok(v) => v,
        Err(e) => {
            let error = BenchmarkError {
                error: format!("Invalid overlap: {}", e),
            };
            eprintln!("{}", serde_json::to_string(&error).unwrap());
            std::process::exit(1);
        }
    };

    let result = run_benchmark(strategy, source_type, path, chunk_size, overlap);

    match result {
        Ok(bench_result) => {
            println!("{}", serde_json::to_string(&bench_result).unwrap());
        }
        Err(e) => {
            let error = BenchmarkError {
                error: format!("Benchmark failed: {}", e),
            };
            eprintln!("{}", serde_json::to_string(&error).unwrap());
            std::process::exit(1);
        }
    }
}

fn run_benchmark(
    strategy: &str,
    source_type: &str,
    path: &str,
    chunk_size: usize,
    overlap: usize,
) -> Result<BenchmarkResult, Box<dyn std::error::Error>> {
    // Parse the source based on source_type
    let source = match source_type {
        "file" => Source::File(path.to_string()),
        "http" | "https" => Source::Http(path.to_string()),
        "string" => Source::Text(path.to_string()),
        _ => {
            return Err(format!(
                "Invalid source_type '{}'. Use 'file', 'string', 'http', 'text', or 'glob'",
                source_type
            )
            .into());
        }
    };

    // Create the chunker using ChunkerBuilder
    match strategy {
        "bytes" => {
            let chunker = ChunkerBuilder::by_bytes(chunk_size, overlap)?;
            bench_with(chunker, source)
        }
        "chars" => {
            let chunker = ChunkerBuilder::by_characters(chunk_size, overlap)?;
            bench_with(chunker, source)
        }
        _ => Err(format!("Invalid strategy '{}'. Use 'bytes' or 'chars'", strategy).into()),
    }
}

// Generic benchmarking body specialized for the concrete chunker type.
fn bench_with<C>(
    chunker: kiru::ChunkerWithStrategy<C>,
    source: Source,
) -> Result<BenchmarkResult, Box<dyn std::error::Error>>
where
    C: kiru::Chunker,
{
    let start = Instant::now();
    let mut num_chunks = 0usize;
    let mut total_bytes = 0usize;

    let iterator = chunker.on_source(source)?;
    for chunk in iterator {
        num_chunks += 1;
        total_bytes += chunk.len();
        std::hint::black_box(chunk.len());
    }

    let elapsed_secs = start.elapsed().as_secs_f64();
    let throughput_mb_s = (total_bytes as f64) / (1024.0 * 1024.0) / elapsed_secs;

    Ok(BenchmarkResult {
        elapsed_secs,
        num_chunks,
        total_bytes,
        throughput_mb_s,
    })
}