Skip to content

Commit c20f9f5

Browse files
Add ColSmol embedding functionality and associated example
- Introduced a new `colsmol_ort` module for ColSmol embedding, including the `OrtColSmolEmbedder` struct with methods for embedding images, files, and queries. - Added a new example in `colsmol.rs` demonstrating how to use the ColSmol embedder with command-line argument parsing. - Updated the `mod.rs` file to include the new `colsmol_ort` module under the `embeddings/local` directory. - Enhanced the `Idefics3ImageProcessor` with additional preprocessing capabilities for better integration with the ColSmol model. - Implemented tests for the new embedding functionality to ensure reliability and correctness.
1 parent 59a7e2d commit c20f9f5

4 files changed

Lines changed: 763 additions & 75 deletions

File tree

rust/examples/colsmol.rs

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
use clap::{Parser, ValueEnum};
2+
3+
#[cfg(feature = "ort")]
4+
use embed_anything::embeddings::local::colsmol_ort::{OrtColSmolEmbedder, ColSmolEmbed};
5+
6+
7+
8+
#[derive(Parser, Debug, Clone, ValueEnum)]
9+
enum ModelType {
10+
Ort,
11+
}
12+
13+
#[derive(Parser, Debug)]
14+
#[command(author, version, about, long_about = None)]
15+
struct Args {
16+
/// Choose model type: 'ort' or 'normal'
17+
#[arg(short, long, default_value = "normal")]
18+
model_type: ModelType,
19+
}
20+
21+
fn main() -> Result<(), anyhow::Error> {
22+
let args = Args::parse();
23+
24+
let colpali_model = match args.model_type {
25+
ModelType::Ort => {
26+
#[cfg(feature = "ort")]
27+
{
28+
Box::new(OrtColSmolEmbedder::new(
29+
"onnx-community/colSmol-256M-ONNX",
30+
None,
31+
Some("onnx"),
32+
)?) as Box<dyn ColSmolEmbed>
33+
}
34+
#[cfg(not(feature = "ort"))]
35+
{
36+
panic!("ORT is not supported without ORT");
37+
}
38+
}
39+
40+
};
41+
// ... rest of the code ...
42+
43+
let image_path = "test.jpg";
44+
let image = image::open(image_path)?;
45+
let embed_data = colpali_model.embed_image(image_path.into(), None)?;
46+
println!("{:?}", embed_data);
47+
// let file_path = "test_files/colpali.pdf";
48+
// let batch_size = 4;
49+
// let embed_data = colpali_model.embed_file(file_path.into(), batch_size)?;
50+
// println!("{:?}", embed_data.len());
51+
52+
// let prompt = "What is attention?";
53+
// let query_embeddings = colpali_model.embed_query(prompt)?;
54+
// println!("{:?}", query_embeddings.len());
55+
Ok(())
56+
}

0 commit comments

Comments
 (0)