Skip to content

Commit

Permalink
feat(benchmarks): add benchmark for simple local pipeline
Browse files Browse the repository at this point in the history
  • Loading branch information
timonv committed Jun 23, 2024
1 parent 477a284 commit 1567940
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 1 deletion.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 6 additions & 1 deletion benchmarks/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,16 @@ tokio = { version = "1.0", features = ["full"] }
swiftide = { path = "../swiftide/", features = ["all"] }
tracing-subscriber = "0.3"
serde_json = "1.0"
criterion = { version = "0.5.1", features = ["html_reports"] }
criterion = { version = "0.5.1", features = ["html_reports", "async_tokio"] }
anyhow = "1.0"
futures-util = "0.3"

[[bench]]
name = "fileloader"
path = "fileloader.rs"
harness = false

[[bench]]
name = "local_pipeline"
path = "local_pipeline.rs"
harness = false
30 changes: 30 additions & 0 deletions benchmarks/local_pipeline.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
use anyhow::Result;
use criterion::{criterion_group, criterion_main, Criterion};
use swiftide::{
ingestion::IngestionPipeline,
integrations::{fastembed::FastEmbed, qdrant::Qdrant},
loaders::FileLoader,
transformers::Embed,
};

async fn run_pipeline() -> Result<()> {
let qdrant_url = "http://localhost:6334";
IngestionPipeline::from_loader(FileLoader::new(".").with_extensions(&["rs"]))
.then_in_batch(10, Embed::new(FastEmbed::builder().batch_size(10).build()?))
.then_store_with(
Qdrant::try_from_url(qdrant_url)?
.batch_size(50)
.vector_size(384)
.collection_name("swiftide-examples-fastembed".to_string())
.build()?,
)
.run()
.await
}

fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("run_local_pipeline", |b| b.iter(run_pipeline));
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);

0 comments on commit 1567940

Please sign in to comment.