Idefics3

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")

processors/src/docx_processor.rs

@@ -1,8 +1,8 @@
-use std::path::Path;
-use docx_parser::MarkdownDocument;
-use text_splitter::ChunkConfigError;
 use crate::markdown_processor::MarkdownProcessor;
 use crate::processor::{Document, DocumentProcessor, FileProcessor};
+use docx_parser::MarkdownDocument;
+use std::path::Path;
+use text_splitter::ChunkConfigError;
 
 /// A struct for processing PDF files.
 pub struct DocxProcessor {
@@ -12,9 +12,7 @@ pub struct DocxProcessor {
 impl DocxProcessor {
     pub fn new(chunk_size: usize, overlap: usize) -> Result<DocxProcessor, ChunkConfigError> {
         let markdown_processor = MarkdownProcessor::new(chunk_size, overlap)?;
-        Ok(DocxProcessor {
-            markdown_processor,
-        })
+        Ok(DocxProcessor { markdown_processor })
     }
 }
 
@@ -34,16 +32,20 @@ mod tests {
         let txt_file = "../test_files/test.docx";
         let processor = DocxProcessor::new(128, 0).unwrap();
 
-        let text = processor.process_file(&txt_file).unwrap();
-        assert!(text.chunks.contains(&"This is a docx file test".to_string()));
+        let text = processor.process_file(txt_file).unwrap();
+        assert!(text
+            .chunks
+            .contains(&"This is a docx file test".to_string()));
     }
 
     // Returns an error if the file path is invalid.
     #[test]
-    #[should_panic(expected = "Error processing file: IO(Os { code: 2, kind: NotFound, message: \"No such file or directory\" })")]
+    #[should_panic(
+        expected = "Error processing file: IO(Os { code: 2, kind: NotFound, message: \"No such file or directory\" })"
+    )]
     fn test_extract_text_invalid_file_path() {
         let invalid_file_path = "this_file_definitely_does_not_exist.docx";
         let processor = DocxProcessor::new(128, 0).unwrap();
-        processor.process_file(&invalid_file_path).unwrap();
+        processor.process_file(invalid_file_path).unwrap();
     }
 }

processors/src/html_processor.rs

@@ -1,8 +1,8 @@
+use crate::markdown_processor::MarkdownProcessor;
+use crate::processor::{Document, DocumentProcessor};
 use anyhow::Result;
 use htmd::{HtmlToMarkdown, HtmlToMarkdownBuilder};
 use text_splitter::ChunkConfigError;
-use crate::markdown_processor::MarkdownProcessor;
-use crate::processor::{Document, DocumentProcessor};
 
 pub struct HtmlDocument {
     pub content: String,
@@ -18,8 +18,7 @@ pub struct HtmlProcessor {
 impl HtmlProcessor {
     pub fn new(chunk_size: usize, overlap: usize) -> Result<HtmlProcessor, ChunkConfigError> {
         let markdown_processor = MarkdownProcessor::new(chunk_size, overlap)?;
-        let html_to_markdown = HtmlToMarkdownBuilder::new()
-            .build();
+        let html_to_markdown = HtmlToMarkdownBuilder::new().build();
         Ok(HtmlProcessor {
             markdown_processor,
             html_to_markdown,
@@ -36,8 +35,8 @@ impl DocumentProcessor for HtmlProcessor {
 
 #[cfg(test)]
 mod tests {
-    use crate::processor::FileProcessor;
     use super::*;
+    use crate::processor::FileProcessor;
 
     #[test]
     fn test_process_html_file() {

processors/src/markdown_processor.rs

@@ -1,30 +1,26 @@
-use text_splitter::{Characters, ChunkConfig, ChunkConfigError, MarkdownSplitter};
 use crate::processor::{Document, DocumentProcessor};
+use text_splitter::{Characters, ChunkConfig, ChunkConfigError, MarkdownSplitter};
 
 /// A struct that provides functionality to process Markdown files.
 pub struct MarkdownProcessor {
-    splitter: MarkdownSplitter<Characters>
+    splitter: MarkdownSplitter<Characters>,
 }
 
 impl MarkdownProcessor {
     pub fn new(chunk_size: usize, overlap: usize) -> Result<MarkdownProcessor, ChunkConfigError> {
-        let splitter_config = ChunkConfig::new(chunk_size)
-            .with_overlap(overlap)?;
+        let splitter_config = ChunkConfig::new(chunk_size).with_overlap(overlap)?;
         let splitter = MarkdownSplitter::new(splitter_config);
-        Ok(MarkdownProcessor {
-            splitter
-        })
+        Ok(MarkdownProcessor { splitter })
     }
 }
 
 impl DocumentProcessor for MarkdownProcessor {
-
     fn process_document(&self, content: &str) -> anyhow::Result<Document> {
-        let chunks = self.splitter.chunks(content)
+        let chunks = self
+            .splitter
+            .chunks(content)
             .map(|x| x.to_string())
             .collect();
-        Ok(Document {
-            chunks
-        })
+        Ok(Document { chunks })
     }
 }

processors/src/pdf/tesseract/command.rs

@@ -4,9 +4,9 @@ use super::*;
 use std::process::{Command, Stdio};
 use std::string::ToString;
 
+use crate::pdf::tesseract::error::{TessError, TessResult};
 #[cfg(target_os = "windows")]
 use std::os::windows::process::CommandExt;
-use crate::pdf::tesseract::error::{TessError, TessResult};
 
 #[cfg(target_os = "windows")]
 const CREATE_NO_WINDOW: u32 = 0x08000000;

processors/src/pdf/tesseract/input.rs

@@ -1,10 +1,10 @@
+use crate::pdf::tesseract::error::{TessError, TessResult};
 use image::DynamicImage;
 use std::{
     collections::HashMap,
     fmt::{self},
     path::{Path, PathBuf},
 };
-use crate::pdf::tesseract::error::{TessError, TessResult};
 
 #[derive(Clone, Debug, PartialEq)]
 pub struct Args {
@@ -121,7 +121,10 @@ mod tests {
     fn test_from_path() {
         let input = Image::from_path("../test_files/clip/cat1.jpg").unwrap();
 
-        assert_eq!(input.get_image_path().unwrap(), "../test_files/clip/cat1.jpg")
+        assert_eq!(
+            input.get_image_path().unwrap(),
+            "../test_files/clip/cat1.jpg"
+        )
     }
 
     #[test]

processors/src/pdf/tesseract/output_boxes.rs

@@ -1,7 +1,7 @@
-use core::fmt;
 use crate::pdf::tesseract::error::TessResult;
 use crate::pdf::tesseract::input::{Args, Image};
 use crate::pdf::tesseract::parse_line_util::{parse_next, FromLine};
+use core::fmt;
 
 #[derive(Debug, PartialEq)]
 pub struct BoxOutput {

processors/src/pdf/tesseract/output_config_parameters.rs

@@ -45,8 +45,7 @@ impl FromLine for ConfigParameter {
     }
 }
 
-pub fn get_tesseract_config_parameters(
-) -> error::TessResult<ConfigParameterOutput> {
+pub fn get_tesseract_config_parameters() -> error::TessResult<ConfigParameterOutput> {
     let mut command = command::get_tesseract_command(None);
     command.arg("--print-parameters");
 
@@ -72,7 +71,9 @@ fn string_to_config_parameter_output(
 
 #[cfg(test)]
 mod tests {
-    use crate::pdf::tesseract::output_config_parameters::{string_to_config_parameter_output, ConfigParameter};
+    use crate::pdf::tesseract::output_config_parameters::{
+        string_to_config_parameter_output, ConfigParameter,
+    };
 
     #[test]
     fn test_string_to_config_parameter_output() {
@@ -97,7 +98,8 @@ mod tests {
     #[test]
     fn test_get_tesseract_config_parameters() {
         let result =
-            crate::pdf::tesseract::output_config_parameters::get_tesseract_config_parameters().unwrap();
+            crate::pdf::tesseract::output_config_parameters::get_tesseract_config_parameters()
+                .unwrap();
         let x = result
             .config_parameters
             .iter()

processors/src/pdf/tesseract/output_data.rs

@@ -73,10 +73,7 @@ impl FromLine for Data {
     }
 }
 
-pub fn image_to_data(
-    image: &Image,
-    args: &Args,
-) -> error::TessResult<DataOutput> {
+pub fn image_to_data(image: &Image, args: &Args) -> error::TessResult<DataOutput> {
     let mut command = command::create_tesseract_command(image, args)?;
     command.arg("tsv");
 
@@ -112,13 +109,12 @@ mod tests {
                 top: 41,
                 width: 46,
                 height: 20,
-                conf: 96.063751,
+                conf: 96.063_75,
                 text: String::from("The"),
             }
         )
     }
 
-
     #[test]
     fn test_string_to_data_parse_error() {
         let result = string_to_data("level   page_num        block_num       par_num line_num        word_num        left    top     width   height  conf    text\n\

processors/src/processor.rs

@@ -1,34 +1,32 @@
 use std::path::Path;
 
 pub trait DocumentProcessor {
-
     fn process_document(&self, content: &str) -> anyhow::Result<Document>;
 }
 
 pub trait FileProcessor {
-
     fn process_file(&self, path: impl AsRef<Path>) -> anyhow::Result<Document>;
 }
 
 pub trait UrlProcessor {
     fn process_url(&self, url: &str) -> anyhow::Result<Document>;
 }
 
-impl <T: DocumentProcessor> FileProcessor for T {
+impl<T: DocumentProcessor> FileProcessor for T {
     fn process_file(&self, path: impl AsRef<Path>) -> anyhow::Result<Document> {
         let bytes = std::fs::read(path)?;
         let out = String::from_utf8_lossy(&bytes);
         self.process_document(&out)
     }
 }
 
-impl <T: DocumentProcessor> UrlProcessor for T {
+impl<T: DocumentProcessor> UrlProcessor for T {
     fn process_url(&self, url: &str) -> anyhow::Result<Document> {
         let content = reqwest::blocking::get(url)?.text()?;
         self.process_document(&content)
     }
 }
 
 pub struct Document {
-    pub chunks: Vec<String>
+    pub chunks: Vec<String>,
 }

processors/src/txt_processor.rs

@@ -1,6 +1,6 @@
-use text_splitter::ChunkConfigError;
 use crate::markdown_processor::MarkdownProcessor;
 use crate::processor::{Document, DocumentProcessor};
+use text_splitter::ChunkConfigError;
 
 /// A struct for processing PDF files.
 pub struct TxtProcessor {
@@ -10,9 +10,7 @@ pub struct TxtProcessor {
 impl TxtProcessor {
     pub fn new(chunk_size: usize, overlap: usize) -> Result<TxtProcessor, ChunkConfigError> {
         let markdown_processor = MarkdownProcessor::new(chunk_size, overlap)?;
-        Ok(TxtProcessor {
-            markdown_processor,
-        })
+        Ok(TxtProcessor { markdown_processor })
     }
 }
 

python/src/config.rs

@@ -25,7 +25,6 @@ impl TextEmbedConfig {
         tesseract_path: Option<&str>,
         pdf_backend: Option<&str>,
     ) -> Self {
-
         let strategy = match splitting_strategy {
             Some(strategy) => {
                 match strategy {

python/src/lib.rs

@@ -212,14 +212,14 @@ impl EmbeddingModel {
             }
             WhichModel::Clip => {
                 let model_id = model_id.unwrap_or("openai/clip-vit-base-patch32");
-                let model = Embedder::Vision(VisionEmbedder::Clip(
+                let model = Embedder::Vision(Box::new(VisionEmbedder::Clip(Box::new(
                     embed_anything::embeddings::local::clip::ClipEmbedder::new(
                         model_id.to_string(),
                         revision,
                         token,
                     )
                     .map_err(|e| PyValueError::new_err(e.to_string()))?,
-                ));
+                ))));
                 Ok(EmbeddingModel {
                     inner: Arc::new(model),
                 })
@@ -264,20 +264,20 @@ impl EmbeddingModel {
                         token,
                         dtype,
                     )
-                    .unwrap(),
+                    .map_err(|e| PyValueError::new_err(e.to_string()))?,
                 )));
                 Ok(EmbeddingModel {
                     inner: Arc::new(model),
                 })
             }
             WhichModel::Colpali => {
                 let model_id = model_id.unwrap_or("vidore/colpali-v1.2-merged");
-                let model = Embedder::Vision(VisionEmbedder::ColPali(Box::new(
+                let model = Embedder::Vision(Box::new(VisionEmbedder::ColPali(Box::new(
                     embed_anything::embeddings::local::colpali::ColPaliEmbedder::new(
                         model_id, revision,
                     )
-                    .unwrap(),
-                )));
+                    .map_err(|e| PyValueError::new_err(e.to_string()))?,
+                ))));
                 Ok(EmbeddingModel {
                     inner: Arc::new(model),
                 })
@@ -321,12 +321,12 @@ impl EmbeddingModel {
             }
             WhichModel::CohereVision => {
                 let model_id = model_id.unwrap_or("embed-v4.0");
-                let model = Embedder::Vision(VisionEmbedder::Cohere(
+                let model = Embedder::Vision(Box::new(VisionEmbedder::Cohere(
                     embed_anything::embeddings::cloud::cohere::CohereEmbedder::new(
                         model_id.to_string(),
                         api_key,
                     ),
-                ));
+                )));
                 Ok(EmbeddingModel {
                     inner: Arc::new(model),
                 })