Enhance Colpali and ColSmol integration with new embedding functionality

- Updated `colpali.py` to enable loading of the Colpali model directly, while commenting out the ONNX model loading for clarity.
- Enhanced the `ColpaliModel` in `colpali.rs` to support both Colpali and ColSmol models based on the model ID provided.
- Introduced a new `ColSmolEmbedder` in `colsmol.rs` for embedding images and files, with improved preprocessing capabilities.
- Updated the `colsmol` module in `mod.rs` to include the new `ColSmolEmbedder`.
- Refactored example usage in `colsmol.rs` to demonstrate the new embedding functionality with command-line argument parsing.
- Added tests for the new embedding features to ensure reliability and correctness.

Author: akshayballal95

examples/colpali.py

@@ -5,12 +5,12 @@
 
 
 # Load the model
-# model: ColpaliModel = ColpaliModel.from_pretrained("vidore/colpali-v1.2-merged", None)
+model: ColpaliModel = ColpaliModel.from_pretrained("vidore/colpali-v1.2-merged", None)
 
 # Load ONNX Model
-model: ColpaliModel = ColpaliModel.from_pretrained_onnx(
-    "starlight-ai/colpali-v1.2-merged-onnx", None
-)
+# model: ColpaliModel = ColpaliModel.from_pretrained_onnx(
+#     "starlight-ai/colpali-v1.2-merged-onnx", None
+# )
 
 # Get all PDF files in the directory
 directory = Path("test_files")

python/src/models/colpali.rs

@@ -1,6 +1,7 @@
 use embed_anything::embeddings::local::colpali::ColPaliEmbed;
 use embed_anything::embeddings::local::colpali::ColPaliEmbedder;
 use embed_anything::embeddings::local::colpali_ort::OrtColPaliEmbedder;
+use embed_anything::embeddings::local::colsmol::ColSmolEmbedder;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::PyResult;
@@ -16,20 +17,40 @@ impl ColpaliModel {
     #[new]
     #[pyo3(signature = (model_id, revision=None))]
     pub fn new(model_id: &str, revision: Option<&str>) -> PyResult<Self> {
-        let model = ColPaliEmbedder::new(model_id, revision)
-            .map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let model_id_small = model_id.to_lowercase();
+        let model = if model_id_small.contains("colpali") {
+            Box::new(
+                ColPaliEmbedder::new(model_id, revision)
+                    .map_err(|e| PyValueError::new_err(e.to_string()))?,
+            ) as Box<dyn ColPaliEmbed + Send + Sync>
+        } else if model_id_small.contains("colsmol") {
+            Box::new(
+                ColSmolEmbedder::new(model_id, revision)
+                    .map_err(|e| PyValueError::new_err(e.to_string()))?,
+            ) as Box<dyn ColPaliEmbed + Send + Sync>
+        } else {
+            return Err(PyValueError::new_err("Invalid model ID"));
+        };
         Ok(Self {
-            model: Box::new(model),
+            model,
         })
     }
 
     #[staticmethod]
     #[pyo3(signature = (model_id, revision=None))]
     pub fn from_pretrained(model_id: &str, revision: Option<&str>) -> PyResult<Self> {
-        let model = ColPaliEmbedder::new(model_id, revision)
-            .map_err(|e| PyValueError::new_err(e.to_string()))?;
+        let model_id_small = model_id.to_lowercase();
+        let model = if model_id_small.contains("colpali") {
+            Box::new(ColPaliEmbedder::new(model_id, revision)
+                .map_err(|e| PyValueError::new_err(e.to_string()))?) as Box<dyn ColPaliEmbed + Send + Sync>
+        } else if model_id_small.contains("colsmol") {
+            Box::new(ColSmolEmbedder::new(model_id, revision)
+                .map_err(|e| PyValueError::new_err(e.to_string()))?) as Box<dyn ColPaliEmbed + Send + Sync>
+        } else {
+            return Err(PyValueError::new_err("Invalid model ID"));
+        };
         Ok(Self {
-            model: Box::new(model),
+            model,
         })
     }
 

rust/examples/colsmol.rs

@@ -1,13 +1,11 @@
 use clap::{Parser, ValueEnum};
 
-#[cfg(feature = "ort")]
-use embed_anything::embeddings::local::colsmol_ort::{OrtColSmolEmbedder, ColSmolEmbed};
-
-
+use embed_anything::embeddings::local::{colpali::ColPaliEmbed, colsmol::ColSmolEmbedder};
 
 #[derive(Parser, Debug, Clone, ValueEnum)]
 enum ModelType {
     Ort,
+    Normal,
 }
 
 #[derive(Parser, Debug)]
@@ -23,34 +21,22 @@ fn main() -> Result<(), anyhow::Error> {
 
     let colpali_model = match args.model_type {
         ModelType::Ort => {
-            #[cfg(feature = "ort")]
-            {
-                Box::new(OrtColSmolEmbedder::new(
-                    "onnx-community/colSmol-256M-ONNX",
-                    None,
-                    Some("onnx"),
-                )?) as Box<dyn ColSmolEmbed>
-            }
-            #[cfg(not(feature = "ort"))]
-            {
-                panic!("ORT is not supported without ORT");
-            }
+            panic!("ORT is not supported without ORT");
+        }
+        ModelType::Normal => {
+            Box::new(ColSmolEmbedder::new("akshayballal/colSmol-256M-merged", None)?)
+                as Box<dyn ColPaliEmbed>
         }
-     
     };
     // ... rest of the code ...
 
-    let image_path = "test.jpg";
-    let image = image::open(image_path)?;
-    let embed_data = colpali_model.embed_image(image_path.into(), None)?;
-    println!("{:?}", embed_data);
-    // let file_path = "test_files/colpali.pdf";
-    // let batch_size = 4;
-    // let embed_data = colpali_model.embed_file(file_path.into(), batch_size)?;
-    // println!("{:?}", embed_data.len());
+    let file_path = "test_files/colpali.pdf";
+    let batch_size = 2;
+    let embed_data = colpali_model.embed_file(file_path.into(), batch_size)?;
+    println!("{:?}", embed_data.len());
 
-    // let prompt = "What is attention?";
-    // let query_embeddings = colpali_model.embed_query(prompt)?;
-    // println!("{:?}", query_embeddings.len());
+    let prompt = "What is attention?";
+    let query_embeddings = colpali_model.embed_query(prompt)?;
+    println!("{:?}", query_embeddings.len());
     Ok(())
 }

rust/src/embeddings/local/colsmol.rs

@@ -0,0 +1,215 @@
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::RwLock;
+
+use base64::Engine;
+use candle_core::{DType, Device};
+use candle_nn::VarBuilder;
+use image::ImageFormat;
+
+use crate::embeddings::embed::{EmbedData, EmbeddingResult};
+use crate::embeddings::local::colpali::{get_images_from_pdf, ColPaliEmbed};
+use crate::embeddings::select_device;
+use crate::models::idefics3::model::{ColIdefics3Model, Idefics3Config};
+use crate::{
+    models::idefics3::tensor_processing::Idefics3Processor,
+};
+
+pub struct ColSmolEmbedder {
+    pub model: RwLock<ColIdefics3Model>,
+    pub processor: Idefics3Processor,
+    pub device: Device,
+}
+
+impl ColSmolEmbedder {
+    pub fn new(model_id: &str, revision: Option<&str>) -> Result<Self, anyhow::Error> {
+        let api = hf_hub::api::sync::Api::new()?;
+        let repo: hf_hub::api::sync::ApiRepo = match revision {
+            Some(rev) => api.repo(hf_hub::Repo::with_revision(
+                model_id.to_string(),
+                hf_hub::RepoType::Model,
+                rev.to_string(),
+            )),
+            None => api.repo(hf_hub::Repo::new(
+                model_id.to_string(),
+                hf_hub::RepoType::Model,
+            )),
+        };
+
+        let model_file = repo.get("model.safetensors")?;
+        let device = select_device();
+        let dtype = if device.is_cuda() {
+            DType::BF16
+        } else {
+            DType::F32
+        };
+
+        let vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)?
+        };
+        let config_file = repo.get("config.json")?;
+
+        let processor = Idefics3Processor::from_pretrained("akshayballal/colSmol-256M-merged")?;
+        let config: Idefics3Config = serde_json::from_slice(&std::fs::read(config_file)?)?;
+
+
+        let model = ColIdefics3Model::load(&config, false, vb)?;
+
+        Ok(Self {
+            model: RwLock::new(model),
+            processor,
+            device,
+        })
+    }
+}
+
+impl ColPaliEmbed for ColSmolEmbedder {
+    fn embed(
+        &self,
+        text_batch: &[&str],
+        batch_size: Option<usize>,
+    ) -> Result<Vec<EmbeddingResult>, anyhow::Error> {
+        let mut encodings = Vec::new();
+        for mini_text_batch in text_batch.chunks(batch_size.unwrap_or(32)) {
+            let (input_ids, attention_mask) = self
+                .processor
+                .tokenize_batch(mini_text_batch.to_vec(), &self.device)?;
+            let batch_encodings = self
+                .model
+                .write()
+                .map_err(|e| anyhow::anyhow!("{}", e))?
+                .forward(&input_ids, &attention_mask, &None, &None)?
+                .to_dtype(DType::F32)?;
+
+            encodings.extend(
+                batch_encodings
+                    .to_vec3::<f32>()?
+                    .iter()
+                    .map(|x| EmbeddingResult::MultiVector(x.to_vec())),
+            );
+        }
+        Ok(encodings)
+    }
+
+    fn embed_query(&self, query: &str) -> anyhow::Result<Vec<EmbedData>> {
+        let (input_ids, attention_mask) =
+            self.processor.tokenize_batch(vec![query], &self.device)?;
+
+        let encoding = self
+            .model
+            .write()
+            .map_err(|e| anyhow::anyhow!("{}", e))?
+            .forward(&input_ids, &attention_mask, &None, &None)?
+            .to_dtype(DType::F32)?
+            .to_vec3::<f32>()?
+            .into_iter()
+            .map(|x| EmbeddingResult::MultiVector(x.to_vec()));
+
+        Ok(encoding
+            .map(|x| EmbedData::new(x.clone(), None, None))
+            .collect::<Vec<_>>())
+    }
+
+    fn embed_image(
+        &self,
+        image_path: PathBuf,
+        metadata: Option<HashMap<String, String>>,
+    ) -> anyhow::Result<EmbedData> {
+        let image = image::open(image_path)?;
+        let (input_ids, attention_mask, pixel_values, pixel_attention_mask) =
+            self.processor.preprocess(&[image], &self.device)?;
+        let encoding = self
+            .model
+            .write()
+            .unwrap()
+            .forward(
+                &input_ids,
+                &attention_mask,
+                &Some(pixel_values),
+                &pixel_attention_mask,
+            )?
+            .to_dtype(DType::F32)?
+            .to_vec3::<f32>()?
+            .into_iter()
+            .map(|x| EmbeddingResult::MultiVector(x.to_vec()))
+            .collect::<Vec<_>>();
+
+        Ok(EmbedData::new(encoding[0].clone(), None, metadata))
+    }
+
+    fn embed_image_batch(&self, image_paths: &[PathBuf]) -> anyhow::Result<Vec<EmbedData>> {
+        let images = image_paths
+            .iter()
+            .map(|path| image::open(path))
+            .collect::<Result<Vec<_>, _>>()?;
+        let (input_ids, attention_mask, pixel_values, pixel_attention_mask) =
+            self.processor.preprocess(&images, &self.device)?;
+        let encodings = self
+            .model
+            .write()
+            .unwrap()
+            .forward(
+                &input_ids,
+                &attention_mask,
+                &Some(pixel_values),
+                &pixel_attention_mask,
+            )?
+            .to_dtype(DType::F32)?
+            .to_vec3::<f32>()?;
+
+        Ok(encodings
+            .into_iter()
+            .map(|x| EmbedData::new(EmbeddingResult::MultiVector(x), None, None))
+            .collect::<Vec<_>>())
+    }
+    fn embed_file(&self, file_path: PathBuf, batch_size: usize) -> anyhow::Result<Vec<EmbedData>> {
+        let pages = get_images_from_pdf(&file_path)?;
+        let mut embed_data = Vec::new();
+        for (index, batch) in pages.chunks(batch_size).enumerate() {
+            let start_page = index * batch_size + 1;
+            let end_page = start_page + batch.len();
+            let page_numbers = (start_page..=end_page).collect::<Vec<_>>();
+            let (input_ids, attention_mask, pixel_values, pixel_attention_mask) =
+                self.processor.preprocess(batch, &self.device)?;
+
+            let image_embeddings = self
+                .model
+                .write()
+                .map_err(|e| anyhow::anyhow!("{}", e))?
+                .forward(
+                    &input_ids,
+                    &attention_mask,
+                    &Some(pixel_values),
+                    &pixel_attention_mask,
+                )?
+                .to_dtype(DType::F32)?
+                .to_vec3::<f32>()?
+                .into_iter()
+                .map(|x| EmbeddingResult::MultiVector(x.to_vec()));
+
+            // zip the embeddings with the page numbers
+            let embed_data_batch = image_embeddings
+                .zip(page_numbers.into_iter())
+                .zip(batch.iter())
+                .map(|((embedding, page_number), page_image)| {
+                    let mut metadata = HashMap::new();
+
+                    let mut buf = Vec::new();
+                    let mut cursor = std::io::Cursor::new(&mut buf);
+                    page_image.write_to(&mut cursor, ImageFormat::Png).unwrap();
+                    let engine = base64::engine::general_purpose::STANDARD;
+                    let base64_image = engine.encode(&buf);
+
+                    metadata.insert("page_number".to_string(), page_number.to_string());
+                    metadata.insert(
+                        "file_path".to_string(),
+                        file_path.to_str().unwrap_or("").to_string(),
+                    );
+                    metadata.insert("image".to_string(), base64_image);
+                    EmbedData::new(embedding, None, Some(metadata))
+                });
+            embed_data.extend(embed_data_batch);
+        }
+        Ok(embed_data)
+    }
+}

rust/src/embeddings/local/mod.rs

@@ -5,8 +5,6 @@ pub mod colbert;
 pub mod colpali;
 #[cfg(feature = "ort")]
 pub mod colpali_ort;
-#[cfg(feature = "ort")]
-pub mod colsmol_ort;
 pub mod jina;
 pub mod model_info;
 pub mod modernbert;
@@ -18,3 +16,4 @@ pub mod pooling;
 pub mod text_embedding;
 pub mod model2vec;
 pub mod qwen3;
+pub mod colsmol;