diff --git a/examples/colpali.py b/examples/colpali.py index 26dc7108..33339765 100644 --- a/examples/colpali.py +++ b/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") diff --git a/processors/src/docx_processor.rs b/processors/src/docx_processor.rs index 52bdbe7b..350da0aa 100644 --- a/processors/src/docx_processor.rs +++ b/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 { 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(); } } diff --git a/processors/src/html_processor.rs b/processors/src/html_processor.rs index 63ce7141..85e0328a 100644 --- a/processors/src/html_processor.rs +++ b/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 { 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() { diff --git a/processors/src/markdown_processor.rs b/processors/src/markdown_processor.rs index a2988dfc..9e1d16ba 100644 --- a/processors/src/markdown_processor.rs +++ b/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 + splitter: MarkdownSplitter, } impl MarkdownProcessor { pub fn new(chunk_size: usize, overlap: usize) -> Result { - 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 { - let chunks = self.splitter.chunks(content) + let chunks = self + .splitter + .chunks(content) .map(|x| x.to_string()) .collect(); - Ok(Document { - chunks - }) + Ok(Document { chunks }) } } diff --git a/processors/src/pdf/tesseract/command.rs b/processors/src/pdf/tesseract/command.rs index b5d003f9..46f0a28c 100644 --- a/processors/src/pdf/tesseract/command.rs +++ b/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; diff --git a/processors/src/pdf/tesseract/input.rs b/processors/src/pdf/tesseract/input.rs index 9773f62f..69865f0f 100644 --- a/processors/src/pdf/tesseract/input.rs +++ b/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] diff --git a/processors/src/pdf/tesseract/output_boxes.rs b/processors/src/pdf/tesseract/output_boxes.rs index 92a1c128..61474337 100644 --- a/processors/src/pdf/tesseract/output_boxes.rs +++ b/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 { diff --git a/processors/src/pdf/tesseract/output_config_parameters.rs b/processors/src/pdf/tesseract/output_config_parameters.rs index d40bca02..f74dedd8 100644 --- a/processors/src/pdf/tesseract/output_config_parameters.rs +++ b/processors/src/pdf/tesseract/output_config_parameters.rs @@ -45,8 +45,7 @@ impl FromLine for ConfigParameter { } } -pub fn get_tesseract_config_parameters( -) -> error::TessResult { +pub fn get_tesseract_config_parameters() -> error::TessResult { 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() diff --git a/processors/src/pdf/tesseract/output_data.rs b/processors/src/pdf/tesseract/output_data.rs index 623153c9..49daefed 100644 --- a/processors/src/pdf/tesseract/output_data.rs +++ b/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 { +pub fn image_to_data(image: &Image, args: &Args) -> error::TessResult { 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\ diff --git a/processors/src/processor.rs b/processors/src/processor.rs index a04d27ce..214d286a 100644 --- a/processors/src/processor.rs +++ b/processors/src/processor.rs @@ -1,12 +1,10 @@ use std::path::Path; pub trait DocumentProcessor { - fn process_document(&self, content: &str) -> anyhow::Result; } pub trait FileProcessor { - fn process_file(&self, path: impl AsRef) -> anyhow::Result; } @@ -14,7 +12,7 @@ pub trait UrlProcessor { fn process_url(&self, url: &str) -> anyhow::Result; } -impl FileProcessor for T { +impl FileProcessor for T { fn process_file(&self, path: impl AsRef) -> anyhow::Result { let bytes = std::fs::read(path)?; let out = String::from_utf8_lossy(&bytes); @@ -22,7 +20,7 @@ impl FileProcessor for T { } } -impl UrlProcessor for T { +impl UrlProcessor for T { fn process_url(&self, url: &str) -> anyhow::Result { let content = reqwest::blocking::get(url)?.text()?; self.process_document(&content) @@ -30,5 +28,5 @@ impl UrlProcessor for T { } pub struct Document { - pub chunks: Vec + pub chunks: Vec, } diff --git a/processors/src/txt_processor.rs b/processors/src/txt_processor.rs index 56348908..cf175eca 100644 --- a/processors/src/txt_processor.rs +++ b/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 { let markdown_processor = MarkdownProcessor::new(chunk_size, overlap)?; - Ok(TxtProcessor { - markdown_processor, - }) + Ok(TxtProcessor { markdown_processor }) } } diff --git a/python/src/config.rs b/python/src/config.rs index 70a981dd..258c4468 100644 --- a/python/src/config.rs +++ b/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 { diff --git a/python/src/lib.rs b/python/src/lib.rs index 50030534..09e4d223 100644 --- a/python/src/lib.rs +++ b/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,7 +264,7 @@ impl EmbeddingModel { token, dtype, ) - .unwrap(), + .map_err(|e| PyValueError::new_err(e.to_string()))?, ))); Ok(EmbeddingModel { inner: Arc::new(model), @@ -272,12 +272,12 @@ impl EmbeddingModel { } 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), }) diff --git a/python/src/models/colpali.rs b/python/src/models/colpali.rs index 0f89038a..179ec19a 100644 --- a/python/src/models/colpali.rs +++ b/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,21 +17,41 @@ impl ColpaliModel { #[new] #[pyo3(signature = (model_id, revision=None))] pub fn new(model_id: &str, revision: Option<&str>) -> PyResult { - let model = ColPaliEmbedder::new(model_id, revision) - .map_err(|e| PyValueError::new_err(e.to_string()))?; - Ok(Self { - model: Box::new(model), - }) + 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 + } 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 + } else { + return Err(PyValueError::new_err("Invalid model ID")); + }; + Ok(Self { model }) } #[staticmethod] #[pyo3(signature = (model_id, revision=None))] pub fn from_pretrained(model_id: &str, revision: Option<&str>) -> PyResult { - let model = ColPaliEmbedder::new(model_id, revision) - .map_err(|e| PyValueError::new_err(e.to_string()))?; - Ok(Self { - model: Box::new(model), - }) + 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 + } 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 + } else { + return Err(PyValueError::new_err("Invalid model ID")); + }; + Ok(Self { model }) } #[staticmethod] diff --git a/rust/examples/bert.rs b/rust/examples/bert.rs index 266f0c1c..2dc9613b 100644 --- a/rust/examples/bert.rs +++ b/rust/examples/bert.rs @@ -1,6 +1,6 @@ use embed_anything::config::{SplittingStrategy, TextEmbedConfig}; -use embed_anything::embeddings::embed::{EmbedData, EmbedderBuilder}; -use embed_anything::{embed_directory_stream, embed_file, embed_files_batch, Dtype}; +use embed_anything::embeddings::embed::EmbedderBuilder; +use embed_anything::{embed_directory_stream, embed_file, embed_files_batch}; use std::collections::HashSet; use std::sync::Arc; use std::{path::PathBuf, time::Instant}; diff --git a/rust/examples/clip.rs b/rust/examples/clip.rs index 900baeaa..97305c17 100644 --- a/rust/examples/clip.rs +++ b/rust/examples/clip.rs @@ -22,7 +22,6 @@ async fn main() { .unwrap() .unwrap(); - let query_emb_data = embed_query(&["Photo of a monkey?"], &model, None) .await .unwrap(); @@ -72,7 +71,7 @@ async fn main() { let mut indices: Vec = (0..similarities.len()).collect(); indices.sort_by(|a, b| similarities[*b].partial_cmp(&similarities[*a]).unwrap()); - + println!("Descending order of similarity: "); for idx in &indices { println!("{}", image_paths[*idx]); diff --git a/rust/examples/cloud.rs b/rust/examples/cloud.rs index e506ab36..3d353476 100644 --- a/rust/examples/cloud.rs +++ b/rust/examples/cloud.rs @@ -1,9 +1,7 @@ use std::{path::PathBuf, sync::Arc}; use embed_anything::{ - config::TextEmbedConfig, - embed_directory_stream, embed_file, - embeddings::embed::Embedder, + config::TextEmbedConfig, embed_directory_stream, embed_file, embeddings::embed::Embedder, }; use anyhow::Result; diff --git a/rust/examples/cohere_pdf.rs b/rust/examples/cohere_pdf.rs index 3c16b6f6..74fbdaf3 100644 --- a/rust/examples/cohere_pdf.rs +++ b/rust/examples/cohere_pdf.rs @@ -1,9 +1,6 @@ use std::sync::Arc; -use embed_anything::{ - config::TextEmbedConfig, - embeddings::embed::Embedder, -}; +use embed_anything::{config::TextEmbedConfig, embeddings::embed::Embedder}; use anyhow::Result; @@ -13,12 +10,13 @@ async fn main() -> Result<()> { .with_chunk_size(1000, Some(0.3)) .with_batch_size(8) .with_buffer_size(8); - let cohere_model: Arc = Arc::new( - Embedder::from_pretrained_cloud("cohere-vision", "embed-v4.0", None).unwrap() - ); + let cohere_model: Arc = + Arc::new(Embedder::from_pretrained_cloud("cohere-vision", "embed-v4.0", None).unwrap()); // let embeddings = cohere_model.embed_query(&["What are Positional Encodings"], Some(&text_embed_config)).await?; - let embeddings = cohere_model.embed_file("test_files/colpali.pdf", Some(&text_embed_config), None).await?; + let embeddings = cohere_model + .embed_file("test_files/colpali.pdf", Some(&text_embed_config), None) + .await?; println!("{:?}", embeddings.unwrap().len()); Ok(()) } diff --git a/rust/examples/colsmol.rs b/rust/examples/colsmol.rs new file mode 100644 index 00000000..19768edd --- /dev/null +++ b/rust/examples/colsmol.rs @@ -0,0 +1,42 @@ +use clap::{Parser, ValueEnum}; + +use embed_anything::embeddings::local::{colpali::ColPaliEmbed, colsmol::ColSmolEmbedder}; + +#[derive(Parser, Debug, Clone, ValueEnum)] +enum ModelType { + Ort, + Normal, +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + /// Choose model type: 'ort' or 'normal' + #[arg(short, long, default_value = "normal")] + model_type: ModelType, +} + +fn main() -> Result<(), anyhow::Error> { + let args = Args::parse(); + + let colpali_model = match args.model_type { + ModelType::Ort => { + panic!("ORT is not supported without ORT"); + } + ModelType::Normal => Box::new(ColSmolEmbedder::new( + "akshayballal/colSmol-256M-merged", + None, + )?) as Box, + }; + // ... rest of the code ... + + 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()); + Ok(()) +} diff --git a/rust/examples/model2vec.rs b/rust/examples/model2vec.rs index d0790ca3..8587578d 100644 --- a/rust/examples/model2vec.rs +++ b/rust/examples/model2vec.rs @@ -67,11 +67,7 @@ async fn main() { // Embed an html file let _out2 = model - .embed_webpage( - "https://www.google.com".to_string(), - Some(&config), - None, - ) + .embed_webpage("https://www.google.com".to_string(), Some(&config), None) .await .unwrap() .unwrap(); diff --git a/rust/src/chunkers/statistical.rs b/rust/src/chunkers/statistical.rs index 243591d4..a3273efb 100644 --- a/rust/src/chunkers/statistical.rs +++ b/rust/src/chunkers/statistical.rs @@ -347,16 +347,16 @@ impl StatisticalChunker { #[cfg(test)] mod tests { - use std::path::PathBuf; - use processors_rs::pdf::pdf_processor::{OcrConfig, PdfBackend}; use crate::extract_document; + use processors_rs::pdf::pdf_processor::{OcrConfig, PdfBackend}; + use std::path::PathBuf; use super::*; #[tokio::test] async fn test_statistical_chunker() { let text = extract_document( - &PathBuf::from("../test_files/attention.pdf"), + PathBuf::from("../test_files/attention.pdf"), 10, 0, OcrConfig { @@ -371,6 +371,6 @@ mod tests { ..Default::default() }; let chunks = chunker.chunk(&text.chunks.join("\n"), 10).await; - assert!(chunks.len() > 0); + assert!(!chunks.is_empty()); } } diff --git a/rust/src/embeddings/cloud/cohere.rs b/rust/src/embeddings/cloud/cohere.rs index bed92602..240689a2 100644 --- a/rust/src/embeddings/cloud/cohere.rs +++ b/rust/src/embeddings/cloud/cohere.rs @@ -88,7 +88,6 @@ impl CohereEmbedder { } pub async fn embed(&self, text_batch: &[&str]) -> Result, anyhow::Error> { - let response = self .client .post(&self.url) @@ -104,7 +103,6 @@ impl CohereEmbedder { .send() .await?; - let data = match response.error_for_status() { Ok(resp) => resp.json::().await?, Err(e) => { @@ -191,9 +189,7 @@ impl CohereEmbedder { } }; let encodings = data.embeddings; - let embedding = encodings - .float - .iter().cloned(); + let embedding = encodings.float.iter().cloned(); embeddings.extend(embedding); } let embeddings = embeddings @@ -321,5 +317,4 @@ mod tests { let embeddings = cohere.embed_pdf(file_path, None).await.unwrap(); assert_eq!(embeddings.len(), 26); } - } diff --git a/rust/src/embeddings/embed.rs b/rust/src/embeddings/embed.rs index 9c13d453..2247d55f 100644 --- a/rust/src/embeddings/embed.rs +++ b/rust/src/embeddings/embed.rs @@ -159,9 +159,9 @@ impl TextEmbedder { token, )?))) } - "model2vec" | "Model2Vec" | "MODEL2VEC" => Ok(Self::Model2Vec(Box::new(Model2VecEmbedder::new( - model_id, token, None, - )?))), + "model2vec" | "Model2Vec" | "MODEL2VEC" => Ok(Self::Model2Vec(Box::new( + Model2VecEmbedder::new(model_id, token, None)?, + ))), "modernbert" | "ModernBert" | "MODERNBERT" => { Ok(Self::ModernBert(Box::new(ModernBertEmbedder::new( @@ -175,7 +175,7 @@ impl TextEmbedder { model_id, revision.map(|s| s.to_string()), token, - dtype + dtype, )?))), _ => Err(anyhow::anyhow!("Model not supported")), } @@ -245,15 +245,15 @@ impl TextEmbedder { /// # Arguments /// /// * `model` - A string holds the model to be used for embedding. Choose from - /// - "openai" - /// - "cohere" + /// - "openai" + /// - "cohere" /// /// * `model_id` - A string holds the model ID for the model to be used for embedding. /// - For OpenAI, find available models at /// - For Cohere, find available models at /// * `api_key` - An optional string holds the API key for authenticating requests to the Cohere API. If not provided, it is taken from the environment variable - /// - For OpenAI, create environment variable `OPENAI_API_KEY` - /// - For Cohere, create environment variable `CO_API_KEY` + /// - For OpenAI, create environment variable `OPENAI_API_KEY` + /// - For Cohere, create environment variable `CO_API_KEY` /// /// # Returns /// @@ -278,21 +278,21 @@ impl TextEmbedder { } pub enum VisionEmbedder { - Clip(ClipEmbedder), + Clip(Box), ColPali(Box), Cohere(CohereEmbedder), } impl From for Embedder { fn from(value: VisionEmbedder) -> Self { - Embedder::Vision(value) + Embedder::Vision(Box::new(value)) } } impl From for VisionEmbedder { fn from(value: Embedder) -> Self { match value { - Embedder::Vision(value) => value, + Embedder::Vision(value) => *value, _ => panic!("Invalid embedder type"), } } @@ -315,11 +315,11 @@ impl VisionEmbedder { token: Option<&str>, ) -> Result { match model { - "clip" | "Clip" | "CLIP" => Ok(Self::Clip(ClipEmbedder::new( + "clip" | "Clip" | "CLIP" => Ok(Self::Clip(Box::new(ClipEmbedder::new( model_id.to_string(), revision, token, - )?)), + )?))), "colpali" | "ColPali" | "COLPALI" => Ok(Self::ColPali(Box::new(ColPaliEmbedder::new( model_id, revision, )?))), @@ -515,7 +515,7 @@ impl EmbedderBuilder { pub enum Embedder { Text(TextEmbedder), - Vision(VisionEmbedder), + Vision(Box), } impl Embedder { @@ -539,15 +539,12 @@ impl Embedder { dtype: Option, ) -> Result { match model_architecture { - "clip" | "Clip" | "CLIP" => Ok(Self::Vision(VisionEmbedder::from_pretrained_hf( - model_architecture, - model_id, - revision, - token, - )?)), - "colpali" | "ColPali" | "COLPALI" => Ok(Self::Vision( + "clip" | "Clip" | "CLIP" => Ok(Self::Vision(Box::new( VisionEmbedder::from_pretrained_hf(model_architecture, model_id, revision, token)?, - )), + ))), + "colpali" | "ColPali" | "COLPALI" => Ok(Self::Vision(Box::new( + VisionEmbedder::from_pretrained_hf(model_architecture, model_id, revision, token)?, + ))), "bert" | "Bert" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( model_architecture, model_id, @@ -562,14 +559,7 @@ impl Embedder { token, dtype, )?)), - "model2vec" | "Model2Vec" | "MODEL2VEC" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, - model_id, - revision, - token, - dtype, - )?)), - "sparse-bert" | "SparseBert" | "SPARSE-BERT" => { + "model2vec" | "Model2Vec" | "MODEL2VEC" => { Ok(Self::Text(TextEmbedder::from_pretrained_hf( model_architecture, model_id, @@ -578,7 +568,7 @@ impl Embedder { dtype, )?)) } - "modernbert" | "ModernBert" | "MODERNBERT" => { + "sparse-bert" | "SparseBert" | "SPARSE-BERT" => { Ok(Self::Text(TextEmbedder::from_pretrained_hf( model_architecture, model_id, @@ -586,8 +576,8 @@ impl Embedder { token, dtype, )?)) - }, - "qwen3" | "Qwen3" | "QWEN3" => { + } + "modernbert" | "ModernBert" | "MODERNBERT" => { Ok(Self::Text(TextEmbedder::from_pretrained_hf( model_architecture, model_id, @@ -596,6 +586,13 @@ impl Embedder { dtype, )?)) } + "qwen3" | "Qwen3" | "QWEN3" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + model_architecture, + model_id, + revision, + token, + dtype, + )?)), _ => Err(anyhow::anyhow!("Model not supported")), } } @@ -612,9 +609,9 @@ impl Embedder { "cohere" | "Cohere" => Ok(Self::Text(TextEmbedder::from_pretrained_cloud( model, model_id, api_key, )?)), - "cohere-vision" | "CohereVision" => Ok(Self::Vision( + "cohere-vision" | "CohereVision" => Ok(Self::Vision(Box::new( VisionEmbedder::from_pretrained_cloud(model, model_id, api_key)?, - )), + ))), _ => Err(anyhow::anyhow!("Model not supported")), } } diff --git a/rust/src/embeddings/local/bert.rs b/rust/src/embeddings/local/bert.rs index 97a355d3..cf8d29c6 100644 --- a/rust/src/embeddings/local/bert.rs +++ b/rust/src/embeddings/local/bert.rs @@ -422,12 +422,12 @@ mod tests { .embed(&["Hello, world!", "I am a rust programmer"], Some(32)) .unwrap(); let test_embeddings: Vec = vec![ - -3.81771438e-02, - 3.29110473e-02, - -5.45941433e-03, - 1.43699292e-02, - -4.02910188e-02, - -1.16532497e-01, + -3.817_714_4e-2, + 3.291_104_7e-2, + -5.459_414_3e-3, + 1.436_992_9e-2, + -4.029_102e-2, + -1.165_325e-1, ]; let embeddings = embeddings[0].to_dense().unwrap()[0..6].to_vec(); println!("{:?}", embeddings); diff --git a/rust/src/embeddings/local/clip.rs b/rust/src/embeddings/local/clip.rs index 21f4b21f..84775731 100644 --- a/rust/src/embeddings/local/clip.rs +++ b/rust/src/embeddings/local/clip.rs @@ -139,7 +139,7 @@ impl ClipEmbedder { let config: Config = serde_json::from_str(&config_str)?; ( VisionModel::Siglip(siglip::Model::new(&config, vb)?), - config.text_config.max_position_embeddings, + config.text_config.max_position_embeddings, config.text_config.pad_token_id as u32, ) }; @@ -267,9 +267,7 @@ impl ClipEmbedder { .tokenize_sequences(Some(mini_text_batch), &self.tokenizer) .unwrap(); - let batch_encodings = self - .model - .get_text_features(&input_ids)?; + let batch_encodings = self.model.get_text_features(&input_ids)?; let normalized_encodings = div_l2_norm(&batch_encodings)?.to_vec2::()?; encodings.extend( normalized_encodings @@ -295,9 +293,7 @@ impl EmbedImage for ClipEmbedder { let images = self .load_images(image_batch, config.vision_config.image_size) .unwrap(); - let batch_encodings = self - .model - .get_image_features(&images)?; + let batch_encodings = self.model.get_image_features(&images)?; let normalized_encodings = div_l2_norm(&batch_encodings)?.to_vec2::()?; encodings.extend(normalized_encodings); } @@ -337,9 +333,7 @@ impl EmbedImage for ClipEmbedder { .unwrap() .unsqueeze(0) .unwrap(); - let encoding = &self - .model - .get_image_features(&image)?; + let encoding = &self.model.get_image_features(&image)?; let normalized_encoding = &div_l2_norm(encoding)?.to_vec2::()?[0]; Ok(EmbedData::new( EmbeddingResult::DenseVector(normalized_encoding.to_vec()), @@ -353,11 +347,12 @@ impl EmbedImage for ClipEmbedder { _pdf_path: T, _batch_size: Option, ) -> anyhow::Result> { - Err(anyhow::anyhow!("PDF embedding not supported for Clip model")) + Err(anyhow::anyhow!( + "PDF embedding not supported for Clip model" + )) } } - #[cfg(test)] mod tests { use super::*; @@ -411,7 +406,13 @@ mod tests { async fn test_embed_image_batch() { let clip_embedder = ClipEmbedder::default(); let embeddings = clip_embedder - .embed_image_batch(&["../test_files/clip/cat1.jpg", "../test_files/clip/cat2.jpeg"], Some(2)) + .embed_image_batch( + &[ + "../test_files/clip/cat1.jpg", + "../test_files/clip/cat2.jpeg", + ], + Some(2), + ) .await .unwrap(); assert_eq!(embeddings.len(), 2); diff --git a/rust/src/embeddings/local/colbert.rs b/rust/src/embeddings/local/colbert.rs index 1918ad73..77d99794 100644 --- a/rust/src/embeddings/local/colbert.rs +++ b/rust/src/embeddings/local/colbert.rs @@ -8,7 +8,7 @@ use std::{ops::Mul, sync::RwLock}; use anyhow::{Error as E, Result}; use hf_hub::{api::sync::Api, Repo}; -use ndarray::{Array2, Axis}; +use ndarray::{Array2, Axis}; use ort::{ execution_providers::{CUDAExecutionProvider, CoreMLExecutionProvider, ExecutionProvider}, session::{builder::GraphOptimizationLevel, Session}, @@ -84,10 +84,7 @@ impl OrtColbertEmbedder { } }; - let _ = match data_filename { - Ok(data) => Some(data), - Err(_) => None, - }; + let _ = data_filename.ok(); let tokenizer_config = std::fs::read_to_string(tokenizer_config_filename)?; let tokenizer_config: TokenizerConfig = serde_json::from_str(&tokenizer_config)?; @@ -174,7 +171,7 @@ impl ColbertEmbed for OrtColbertEmbedder { let batch_size = batch_size.unwrap_or(32); - let model_guard = self.model.read().unwrap(); + let model_guard = self.model.read().unwrap(); let (input_names, output_name) = { let names = model_guard .inputs @@ -213,8 +210,6 @@ impl ColbertEmbed for OrtColbertEmbedder { mask_row[1] = 1; } } - - let input_ids_tensor = ort::value::TensorRef::from_array_view(&input_ids)?; let attention_mask_tensor = ort::value::TensorRef::from_array_view(&attention_mask)?; let mut inputs = @@ -285,7 +280,6 @@ impl BertEmbed for OrtColbertEmbedder { let (input_ids, attention_mask): (Array2, Array2) = tokenize_batch_ndarray(&self.tokenizer, mini_text_batch)?; let token_type_ids: Array2 = Array2::zeros(input_ids.raw_dim()); - let input_ids_tensor = ort::value::TensorRef::from_array_view(&input_ids)?; let attention_mask_tensor = ort::value::TensorRef::from_array_view(&attention_mask)?; diff --git a/rust/src/embeddings/local/colpali_ort.rs b/rust/src/embeddings/local/colpali_ort.rs index c51cae84..04bcd579 100644 --- a/rust/src/embeddings/local/colpali_ort.rs +++ b/rust/src/embeddings/local/colpali_ort.rs @@ -187,7 +187,6 @@ impl OrtColPaliEmbedder { attention_mask: Array2, pixel_values: Array4, ) -> Result, E> { - let mut model_guard = self.model.write().unwrap(); let output_name = model_guard.outputs.first().unwrap().name.to_string(); @@ -201,7 +200,6 @@ impl OrtColPaliEmbedder { .to_owned() .into_dimensionality::()?; - let e = embeddings .outer_iter() .map(|row| { diff --git a/rust/src/embeddings/local/colsmol.rs b/rust/src/embeddings/local/colsmol.rs new file mode 100644 index 00000000..d6795803 --- /dev/null +++ b/rust/src/embeddings/local/colsmol.rs @@ -0,0 +1,210 @@ +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, + pub processor: Idefics3Processor, + pub device: Device, +} + +impl ColSmolEmbedder { + pub fn new(model_id: &str, revision: Option<&str>) -> Result { + 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, + ) -> Result, 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::()? + .iter() + .map(|x| EmbeddingResult::MultiVector(x.to_vec())), + ); + } + Ok(encodings) + } + + fn embed_query(&self, query: &str) -> anyhow::Result> { + 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::()? + .into_iter() + .map(|x| EmbeddingResult::MultiVector(x.to_vec())); + + Ok(encoding + .map(|x| EmbedData::new(x.clone(), None, None)) + .collect::>()) + } + + fn embed_image( + &self, + image_path: PathBuf, + metadata: Option>, + ) -> anyhow::Result { + 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::()? + .into_iter() + .map(|x| EmbeddingResult::MultiVector(x.to_vec())) + .collect::>(); + + Ok(EmbedData::new(encoding[0].clone(), None, metadata)) + } + + fn embed_image_batch(&self, image_paths: &[PathBuf]) -> anyhow::Result> { + let images = image_paths + .iter() + .map(image::open) + .collect::, _>>()?; + 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::()?; + + Ok(encodings + .into_iter() + .map(|x| EmbedData::new(EmbeddingResult::MultiVector(x), None, None)) + .collect::>()) + } + fn embed_file(&self, file_path: PathBuf, batch_size: usize) -> anyhow::Result> { + 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::>(); + 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::()? + .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) + } +} diff --git a/rust/src/embeddings/local/mod.rs b/rust/src/embeddings/local/mod.rs index 807dc167..43219a27 100644 --- a/rust/src/embeddings/local/mod.rs +++ b/rust/src/embeddings/local/mod.rs @@ -5,7 +5,9 @@ pub mod colbert; pub mod colpali; #[cfg(feature = "ort")] pub mod colpali_ort; +pub mod colsmol; pub mod jina; +pub mod model2vec; pub mod model_info; pub mod modernbert; #[cfg(feature = "ort")] @@ -13,6 +15,5 @@ pub mod ort_bert; #[cfg(feature = "ort")] pub mod ort_jina; pub mod pooling; -pub mod text_embedding; -pub mod model2vec; pub mod qwen3; +pub mod text_embedding; diff --git a/rust/src/embeddings/local/ort_bert.rs b/rust/src/embeddings/local/ort_bert.rs index 65866343..86a8eb6f 100644 --- a/rust/src/embeddings/local/ort_bert.rs +++ b/rust/src/embeddings/local/ort_bert.rs @@ -77,7 +77,10 @@ impl OrtBertEmbedder { let config = api.get("config.json")?; let tokenizer = api.get("tokenizer.json")?; let tokenizer_config = api.get("tokenizer_config.json")?; - let mut base_path = path.rsplit_once('/').map(|(p, _)| p.to_string()).unwrap_or_default(); + let mut base_path = path + .rsplit_once('/') + .map(|(p, _)| p.to_string()) + .unwrap_or_default(); if !base_path.is_empty() { base_path.push('/'); } @@ -182,7 +185,7 @@ impl OrtBertEmbedder { .map(|input| input.name.as_str()) .collect(); let output_name = model_guard.outputs.first().unwrap().name.to_string(); - let needs_token_type = input_names.iter().any(|&x| x == "token_type_ids"); + let needs_token_type = input_names.contains(&"token_type_ids"); // Run model and extract embeddings let encodings = text_batch @@ -259,7 +262,7 @@ impl OrtBertEmbedder { .iter() .map(|input| input.name.as_str()) .collect(); - let needs_token_type = input_names.iter().any(|&x| x == "token_type_ids"); + let needs_token_type = input_names.contains(&"token_type_ids"); let output_name = model_guard.outputs.first().unwrap().name.to_string(); for mini_text_batch in text_batch.chunks(batch_size) { @@ -567,10 +570,10 @@ mod tests { println!("embeddings: {:?}", embeddings); let test_embeddings: Vec = vec![ - -3.81771736e-02, - 3.29111032e-02, - -5.45938499e-03, - 1.43699143e-02, + -3.817_717_4e-2, + 3.291_110_3e-2, + -5.459_385e-3, + 1.436_991_4e-2, ]; let embeddings = embeddings[0].to_dense().unwrap()[0..4].to_vec(); assert!( diff --git a/rust/src/embeddings/local/ort_colsmol.rs b/rust/src/embeddings/local/ort_colsmol.rs new file mode 100644 index 00000000..9d938a4c --- /dev/null +++ b/rust/src/embeddings/local/ort_colsmol.rs @@ -0,0 +1,459 @@ +use std::sync::RwLock; +use std::{collections::HashMap, path::PathBuf}; + +use crate::embeddings::embed::{EmbedData, EmbeddingResult}; +use crate::embeddings::local::colpali::ColPaliEmbed; +use crate::models::idefics3::array_processing::Idefics3Processor; +use crate::models::paligemma; +use anyhow::Error as E; +use base64::Engine; +use half::f16; +use image::ImageFormat; +use ndarray::prelude::*; +use ort::execution_providers::{CUDAExecutionProvider, CoreMLExecutionProvider, ExecutionProvider}; +use ort::session::builder::GraphOptimizationLevel; +use ort::session::Session; +use rayon::prelude::*; +use tokenizers::{PaddingParams, Tokenizer, TruncationParams}; + +use super::colpali::get_images_from_pdf; + + +pub struct OrtColSmolEmbedder { + pub model: RwLock, + pub tokenizer: Tokenizer, + pub image_size: usize, + pub num_channels: usize, + pub processor: Idefics3Processor, + dummy_input: Array2, +} + +impl OrtColSmolEmbedder { + pub fn new( + model_id: &str, + revision: Option<&str>, + path_in_repo: Option<&str>, + ) -> Result { + 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 mut path_in_repo = path_in_repo.unwrap_or_default().to_string(); + if !path_in_repo.is_empty() { + path_in_repo.push('/'); + } + let (_, tokenizer_filename, weights_filename, _, processing_config_filename,) = { + let config = repo + .get("config.json") + .unwrap_or(repo.get("preprocessor_config.json")?); + let tokenizer = repo.get("tokenizer.json")?; + let weights = repo.get(format!("{}model.onnx", path_in_repo).as_str())?; + let data = repo + .get(format!("{}model.onnx_data", path_in_repo).as_str()) + .ok(); + let processing_config = repo.get("preprocessor_config.json")?; + + (config, tokenizer, weights, data, processing_config) + }; + + let config: paligemma::Config = paligemma::Config::paligemma_3b_448(); + let image_size = config.vision_config.image_size; + let num_channels = config.vision_config.num_channels; + + let mut tokenizer = Tokenizer::from_pretrained(model_id, None).map_err(E::msg)?; + + let pp = PaddingParams { + strategy: tokenizers::PaddingStrategy::BatchLongest, + ..Default::default() + }; + let trunc = TruncationParams { + strategy: tokenizers::TruncationStrategy::LongestFirst, + max_length: config.text_config.max_position_embeddings, + ..Default::default() + }; + + let processor: Idefics3Processor = Idefics3Processor::from_pretrained(model_id)?; + + tokenizer + .with_padding(Some(pp)) + .with_truncation(Some(trunc)) + .unwrap(); + + + let cuda = CUDAExecutionProvider::default(); + + if !cuda.is_available()? { + eprintln!("CUDAExecutionProvider is not available"); + } else { + println!("Session is using CUDAExecutionProvider"); + } + + // Get physical core count (excluding hyperthreading) + let threads = std::thread::available_parallelism() + .map(|p| p.get()) + .unwrap_or(1); + // For CPU-bound workloads like ONNX inference, it's often better to use + // physical cores rather than logical cores to avoid context switching overhead + let optimal_threads = std::cmp::max(1, threads / 2); + + let model = Session::builder()? + .with_execution_providers([ + CUDAExecutionProvider::default().build(), + CoreMLExecutionProvider::default().build(), + ])? + .with_optimization_level(GraphOptimizationLevel::Level3)? + .with_intra_threads(optimal_threads)? // Use optimal thread count + .with_inter_threads(1)? // Set inter-op parallelism to 1 when using GPU + .commit_from_file(weights_filename)?; + + let dummy_prompt: &str = "<|im_start|>user\nDescribe the image."; + let dummy_input = tokenize(&tokenizer, dummy_prompt.to_string())?; + + Ok(Self { + model: RwLock::new(model), + tokenizer, + image_size, + processor, + num_channels, + dummy_input, + }) + } +} + +fn tokenize_batch(tokenizer: &Tokenizer, text_batch: &[&str]) -> Result, E> { + let token_ids = tokenizer + .encode_batch_fast(text_batch.to_vec(), true) + .map_err(E::msg)? + .iter() + .map(|tokens| { + tokens + .get_ids() + .iter() + .map(|&id| id as i64) + .collect::>() + }) + .collect::>>(); + + let token_ids_array = Array2::from_shape_vec( + (token_ids.len(), token_ids[0].len()), + token_ids.into_iter().flatten().collect::>(), + ) + .unwrap(); + + Ok(token_ids_array) +} + +fn tokenize(tokenizer: &Tokenizer, text: String) -> Result, E> { + let token_ids = tokenizer.encode(text, false).map_err(E::msg)?; + let token_ids_array = Array2::from_shape_vec( + (1, token_ids.len()), + token_ids + .get_ids() + .iter() + .map(|x| *x as i64) + .collect::>(), + ) + .unwrap(); + Ok(token_ids_array) +} + +fn get_attention_mask(tokenizer: &Tokenizer, text_batch: &[&str]) -> Result, E> { + let attention_mask = tokenizer + .encode_batch(text_batch.to_vec(), true) + .map_err(E::msg)? + .iter() + .map(|tokens| { + tokens + .get_attention_mask() + .to_vec() + .iter() + .map(|x| *x as i64) + .collect::>() + }) + .collect::>>(); + Ok(Array2::from_shape_vec( + (attention_mask.len(), attention_mask[0].len()), + attention_mask.into_iter().flatten().collect::>(), + ) + .unwrap()) +} + +impl OrtColSmolEmbedder { + fn run_model( + &self, + token_ids: Array2, + attention_mask: Array2, + pixel_values: Array5, + pixel_attention_mask: Array4, + ) -> Result, E> { + let mut model_guard = self.model.write().unwrap(); + let output_name = model_guard.outputs.first().unwrap().name.to_string(); + println!("input names: {:?}", model_guard.inputs); + + let token_ids_tensor = ort::value::Value::from_array(token_ids)?; + let pixel_values_tensor = ort::value::Value::from_array(pixel_values)?; + let attention_mask_tensor = ort::value::Value::from_array(attention_mask)?; + let pixel_attention_mask_tensor = ort::value::Value::from_array(pixel_attention_mask)?; + let outputs = model_guard.run(ort::inputs!["input_ids" => token_ids_tensor, "pixel_values" => pixel_values_tensor, "attention_mask" => attention_mask_tensor, "pixel_attention_mask" => pixel_attention_mask_tensor])?; + + let embeddings = outputs[output_name] + .try_extract_array::()? + .to_owned() + .into_dimensionality::()?; + + let e = embeddings + .outer_iter() + .map(|row| { + EmbeddingResult::MultiVector( + row.outer_iter() + .map(|x| x.to_vec()) + .map(|y| y.into_iter().map(|z| z.to_f32()).collect::>()) + .collect(), + ) + }) + .collect(); + + Ok(e) + } +} + +impl ColPaliEmbed for OrtColSmolEmbedder { + fn embed( + &self, + text_batch: &[&str], + batch_size: Option, + ) -> Result, anyhow::Error> { + let batch_size = batch_size.unwrap_or(32); + let encodings = text_batch + .par_chunks(batch_size) + .flat_map(|mini_text_batch| -> Result, E> { + let token_ids: Array2 = tokenize_batch(&self.tokenizer, mini_text_batch)?; + let attention_mask: Array2 = + get_attention_mask(&self.tokenizer, mini_text_batch)?; + let pixel_values: Array5 = Array5::zeros(( + mini_text_batch.len(), + 1, + self.num_channels, + self.image_size, + self.image_size, + )); + let pixel_attention_mask: Array4 = Array4::ones(( + mini_text_batch.len(), + self.num_channels, + self.image_size, + self.image_size, + )); + let e = self.run_model( + token_ids, + attention_mask, + pixel_values, + pixel_attention_mask, + )?; + Ok(e) + }) + .flatten() + .collect::>(); + + Ok(encodings) + } + + fn embed_query(&self, query: &str) -> anyhow::Result> { + let token_ids = tokenize_batch(&self.tokenizer, &[query])?; + let attention_mask = get_attention_mask(&self.tokenizer, &[query])?; + let pixel_values: Array5 = + Array5::zeros((1, 1, self.num_channels, self.image_size, self.image_size)); + let pixel_attention_mask: Array4 = + Array4::ones((1, self.num_channels, self.image_size, self.image_size)); + let e = self + .run_model( + token_ids, + attention_mask, + pixel_values, + pixel_attention_mask, + )? + .into_iter() + .map(|x| EmbedData::new(x, None, None)) + .collect::>(); + Ok(e) + } + + fn embed_file(&self, file_path: PathBuf, batch_size: usize) -> anyhow::Result> { + let pages = get_images_from_pdf(&file_path)?; + let mut embed_data = Vec::new(); + for (index, batch) in pages.chunks(1).enumerate() { + let start_page = index * batch_size + 1; + let end_page = start_page + batch.len(); + let page_numbers = (start_page..=end_page).collect::>(); + let (input_ids, attention_mask, page_images, pixel_attention_mask) = self.processor.preprocess(&batch)?; + + let image_embeddings = self.run_model( + input_ids, + attention_mask, + page_images.insert_axis(Axis(0)), + pixel_attention_mask + .unwrap() + .insert_axis(Axis(0)) + .mapv(|x| x as i64), + )?; + // zip the embeddings with the page numbers + let embed_data_batch = image_embeddings + .into_iter() + .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) + } + + fn embed_image( + &self, + image_path: PathBuf, + metadata: Option>, + ) -> anyhow::Result { + let image = image::open(&image_path)?; + let (input_ids, attention_mask, image_array, pixel_attention_mask) = self.processor.preprocess(&[image])?; + + println!("image array shape: {:?}", image_array.shape()); + println!("input ids {:?}", input_ids); + let e = self + .run_model( + input_ids, + attention_mask, + image_array.insert_axis(Axis(0)), + pixel_attention_mask + .unwrap() + .insert_axis(Axis(0)) + .mapv(|x| x as i64), + )? + .into_iter() + .map(|x| EmbedData::new(x, None, metadata.clone())) + .collect::>(); + Ok(e[0].clone()) + } + + fn embed_image_batch(&self, image_paths: &[PathBuf]) -> anyhow::Result> { + let images = image_paths + .iter() + .map(|path| image::open(path).unwrap()) + .collect::>(); + let (input_ids, attention_mask, image_array, pixel_attention_mask) = self.processor.preprocess(&images)?; + + let e = self + .run_model( + input_ids, + attention_mask, + image_array.insert_axis(Axis(0)), + pixel_attention_mask + .unwrap() + .insert_axis(Axis(0)) + .mapv(|x| x as i64), + )? + .into_iter() + .enumerate() + .map(|(i, x)| { + let mut metadata: HashMap = HashMap::new(); + metadata.insert( + "file_path".to_string(), + image_paths[i].to_str().unwrap().to_string(), + ); + EmbedData::new(x, None, Some(metadata)) + }) + .collect::>(); + Ok(e) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use lazy_static::lazy_static; + use std::fs; + use std::path::Path; + use std::sync::Mutex; + + lazy_static! { + static ref MODEL: Mutex = Mutex::new( + OrtColSmolEmbedder::new("onnx-community/colSmol-256M-ONNX", None, None).unwrap() + ); + } + + const IMAGE_PATH: &str = "temp_image.png"; + const IMAGE_URL: &str = + "https://www.google.com/images/branding/googlelogo/1x/googlelogo_color_272x92dp.png"; + + async fn download_image() -> anyhow::Result<()> { + if !Path::new(IMAGE_PATH).exists() { + let response = reqwest::get(IMAGE_URL).await?; + let bytes = response.bytes().await?; + fs::write(IMAGE_PATH, &bytes)?; + } + Ok(()) + } + + #[tokio::test] + async fn test_colpali_embed_image() -> anyhow::Result<()> { + download_image().await?; + let model = MODEL.lock().unwrap(); + let embedding = model.embed_image(PathBuf::from(IMAGE_PATH), None).unwrap(); + assert!( + !embedding.embedding.to_multi_vector().unwrap().is_empty(), + "Embedding should not be empty" + ); + Ok(()) + } + + #[tokio::test] + async fn test_colpali_embed_image_batch() -> anyhow::Result<()> { + download_image().await?; + let model = MODEL.lock().unwrap(); + let embeddings = model + .embed_image_batch(&[PathBuf::from(IMAGE_PATH), PathBuf::from(IMAGE_PATH)]) + .unwrap(); + assert_eq!(embeddings.len(), 2, "There should be two embeddings"); + Ok(()) + } + + #[tokio::test] + async fn test_colpali_embed_file() -> anyhow::Result<()> { + let model = MODEL.lock().unwrap(); + let embeddings = model + .embed_file(PathBuf::from("../test_files/test.pdf"), 1) + .unwrap(); + assert_eq!(embeddings.len(), 1, "There should be 1 embeddings"); + Ok(()) + } + + #[test] + fn cleanup() -> anyhow::Result<()> { + if Path::new(IMAGE_PATH).exists() { + fs::remove_file(IMAGE_PATH)?; + } + Ok(()) + } +} diff --git a/rust/src/embeddings/local/ort_jina.rs b/rust/src/embeddings/local/ort_jina.rs index 23d00a13..95608619 100644 --- a/rust/src/embeddings/local/ort_jina.rs +++ b/rust/src/embeddings/local/ort_jina.rs @@ -78,7 +78,10 @@ impl OrtJinaEmbedder { let config = api.get("config.json")?; let tokenizer = api.get("tokenizer.json")?; let tokenizer_config = api.get("tokenizer_config.json")?; - let mut base_path = path.rsplit_once('/').map(|(p, _)| p.to_string()).unwrap_or_default(); + let mut base_path = path + .rsplit_once('/') + .map(|(p, _)| p.to_string()) + .unwrap_or_default(); if !base_path.is_empty() { base_path.push('/'); } diff --git a/rust/src/embeddings/local/pooling.rs b/rust/src/embeddings/local/pooling.rs index d8d8163a..ba5e94d8 100644 --- a/rust/src/embeddings/local/pooling.rs +++ b/rust/src/embeddings/local/pooling.rs @@ -88,11 +88,11 @@ impl Pooling { .get_on_dim(1, attention_mask.dim(1)? - 1)? .sum_all()? .to_scalar::()?; - + if left_padding == attention_mask.dim(0)? as u32 { - return Ok(PooledOutputType::Tensor( + Ok(PooledOutputType::Tensor( tensor.get_on_dim(1, tensor.dim(1)? - 1)?, - )); + )) } else { let sequence_lengths = attention_mask.sum(1)?.to_vec1::()?; let batch_size = tensor.dim(0)?; @@ -100,11 +100,11 @@ impl Pooling { // Create a tensor of indices for the last tokens let indices: Vec = sequence_lengths.iter().map(|&len| len - 1).collect(); let indices = Tensor::from_vec(indices, (batch_size,), tensor.device())?; - + // Use gather to get all last tokens at once let final_tensor = tensor.gather(&indices, 1)?; - return Ok(PooledOutputType::Tensor(final_tensor)) - } + Ok(PooledOutputType::Tensor(final_tensor)) + } } ModelOutput::Array(array) => { let attention_mask = attention_mask @@ -117,12 +117,12 @@ impl Pooling { let batch_size = array.shape()[0]; let mut final_embeddings = vec![]; - for i in 0..batch_size{ - let t = array.slice(s![i, .., (sequence_lengths[i]-1.0) as usize]); + for i in 0..batch_size { + let t = array.slice(s![i, .., (sequence_lengths[i] - 1.0) as usize]); final_embeddings.push(t); } let final_embeddings = ndarray::stack(Axis(1), &final_embeddings)?; - return Ok(PooledOutputType::Array(final_embeddings)) + Ok(PooledOutputType::Array(final_embeddings)) } } } diff --git a/rust/src/embeddings/local/qwen3.rs b/rust/src/embeddings/local/qwen3.rs index 56e08cd6..bbce4480 100644 --- a/rust/src/embeddings/local/qwen3.rs +++ b/rust/src/embeddings/local/qwen3.rs @@ -128,7 +128,7 @@ impl Qwen3Embed for Qwen3Embedder { .forward(&token_ids, &attention_mask, 0) .unwrap() .to_dtype(DType::F32)?; - + self.model.write().unwrap().clear_kv_cache(); let attention_mask = PooledOutputType::from(attention_mask); let attention_mask = Some(&attention_mask); diff --git a/rust/src/embeddings/local/text_embedding.rs b/rust/src/embeddings/local/text_embedding.rs index 9620a0fe..a212790a 100644 --- a/rust/src/embeddings/local/text_embedding.rs +++ b/rust/src/embeddings/local/text_embedding.rs @@ -376,7 +376,6 @@ fn init_models_map() -> HashMap> { model_code: String::from("onnx-community/Qwen3-Embedding-0.6B-ONNX"), model_file: String::from("onnx/model.onnx"), }, - ]; // TODO: Use when out in stable diff --git a/rust/src/models/idefics3/array_processing.rs b/rust/src/models/idefics3/array_processing.rs new file mode 100644 index 00000000..ef80dc45 --- /dev/null +++ b/rust/src/models/idefics3/array_processing.rs @@ -0,0 +1,797 @@ +use std::collections::HashMap; + +use anyhow::{Error, Ok}; +use hf_hub::{api::sync::Api, Repo, RepoType}; +use image::{imageops::FilterType, DynamicImage, GenericImageView, RgbImage}; +use ndarray::{s, Array2}; +use regex::Regex; +use serde::Deserialize; +use tokenizers::{AddedToken, PaddingParams, Tokenizer, TruncationParams}; + +const MAX_IMAGE_SIZE: i32 = 4096; + +#[derive(Debug, Clone, Deserialize)] +pub struct Idefics3ImageProcessor { + do_convert_rgb: bool, + do_resize: bool, + size: Option>, + do_image_splitting: bool, + image_mean: Option>, + image_std: Option>, + + #[serde(default = "default_max_image_size")] + max_image_size: Option>, + do_rescale: bool, + rescale_factor: f32, + do_pad: bool, + do_normalize: bool, +} + +fn default_max_image_size() -> Option> { + Some(HashMap::from([("longest_edge".to_string(), 364)])) +} + +impl Idefics3ImageProcessor { + pub fn new( + do_convert_rgb: bool, + do_resize: bool, + size: Option>, + do_image_splitting: bool, + max_image_size: Option>, + do_rescale: bool, + rescale_factor: f32, + do_pad: bool, + do_normalize: bool, + image_mean: Option>, + image_std: Option>, + ) -> Self { + let max_image_size = + max_image_size.unwrap_or(HashMap::from([("longest_edge".to_string(), 364)])); + let image_mean = image_mean.unwrap_or(vec![0.5, 0.5, 0.5]); + let image_std = image_std.unwrap_or(vec![0.5, 0.5, 0.5]); + Self { + do_convert_rgb, + do_resize, + size, + do_image_splitting, + image_mean: Some(image_mean), + image_std: Some(image_std), + max_image_size: Some(max_image_size), + do_rescale, + rescale_factor, + do_pad, + do_normalize, + } + } + + pub fn from_pretrained(model_id: &str) -> Result { + let api = Api::new()?; + let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model)); + let config_file = repo.get("preprocessor_config.json").unwrap(); + let processor: Idefics3ImageProcessor = + serde_json::from_slice(&std::fs::read(config_file).unwrap()).unwrap(); + Ok(processor) + } + + pub fn resize_for_vision_encoder( + &self, + image: &DynamicImage, + vision_encoder_max_size: &i32, + resample: FilterType, + ) -> ndarray::Array3 { + let (width, height) = image.dimensions(); + let height = height as f32; + let width = width as f32; + let vision_encoder_max_size = *vision_encoder_max_size as f32; + + let aspect_ratio = width / height; + let (new_height, new_width) = if width >= height { + let width = (width / vision_encoder_max_size).ceil() * vision_encoder_max_size; + let height = (width / aspect_ratio).floor(); + let height = + ((height / vision_encoder_max_size).ceil() * vision_encoder_max_size).ceil(); + (height, width) + } else { + let height = (height / vision_encoder_max_size).ceil() * vision_encoder_max_size; + let width = (height * aspect_ratio).floor(); + let width = (width / vision_encoder_max_size).ceil() * vision_encoder_max_size; + (height, width) + }; + self.resize( + &image, + HashMap::from([ + ("width".to_string(), new_width as i32), + ("height".to_string(), new_height as i32), + ]), + resample, + ) + } + + pub fn resize( + &self, + image: &DynamicImage, + size: HashMap, + resample: FilterType, + ) -> ndarray::Array3 { + let resized_image = if size.contains_key("height") && size.contains_key("width") { + image.resize_exact( + size.get("width").cloned().unwrap() as u32, + size.get("height").cloned().unwrap() as u32, + resample, + ) + } else { + let size = get_resize_output_image_size( + image.clone(), + size.get("longest_edge").cloned().unwrap(), + ); + image.resize_exact(size.1 as u32, size.0 as u32, resample) + }; + let (width, height) = resized_image.dimensions(); + let resized_image_array = resized_image.to_rgb8().into_raw(); + let resized_image_array = ndarray::Array3::from_shape_vec( + (height as usize, width as usize, 3), + resized_image_array, + ) + .unwrap(); + resized_image_array + } + + pub fn split_image( + &self, + image: &DynamicImage, + max_image_size: &HashMap, + resample: FilterType, + ) -> (Vec>, i32, i32) { + let (width, height) = image.dimensions(); + let max_size = max_image_size.get("longest_edge").unwrap_or(&364); + let max_height = *max_size; + let max_width = *max_size; + + let mut frames = Vec::new(); + let (num_splits_h, num_splits_w) = if height > max_height as u32 || width > max_width as u32 + { + // Calculate the number of splits + let num_splits_h = (height as f32 / max_height as f32).ceil() as i32; + let num_splits_w = (width as f32 / max_width as f32).ceil() as i32; + + // Calculate optimal dimensions for sub-images + let optimal_height = (height as f32 / num_splits_h as f32).ceil() as u32; + let optimal_width = (width as f32 / num_splits_w as f32).ceil() as u32; + + // Iterate through each row and column + for r in 0..num_splits_h { + for c in 0..num_splits_w { + // Calculate crop coordinates + let start_x = (c as u32) * optimal_width; + let start_y = (r as u32) * optimal_height; + let end_x = std::cmp::min(start_x + optimal_width, width); + let end_y = std::cmp::min(start_y + optimal_height, height); + + // Crop the image + let cropped_image = + image.crop_imm(start_x, start_y, end_x - start_x, end_y - start_y); + frames.push(cropped_image); + } + } + + // For the global image at the end, we resize it to match the max_image_size, for cpu memory efficiency + let global_image_height = max_height as u32; + let global_image_width = max_width as u32; + let global_image = if height != global_image_height || width != global_image_width { + let size = HashMap::from([ + ("height".to_string(), global_image_height as i32), + ("width".to_string(), global_image_width as i32), + ]); + let resized = self.resize(image, size, resample); + DynamicImage::ImageRgb8( + RgbImage::from_raw( + resized.shape()[1] as u32, + resized.shape()[0] as u32, + resized.into_raw_vec_and_offset().0, + ) + .unwrap(), + ) + } else { + image.clone() + }; + frames.push(global_image); + + (num_splits_h, num_splits_w) + } else { + // If image is smaller than max_size, just add it as is + frames.push(image.clone()); + (0, 0) + }; + + let frames = frames + .iter() + .map(|frame| { + let image = frame.to_rgb8().into_raw(); + let image = ndarray::Array3::from_shape_vec( + (frame.height() as usize, frame.width() as usize, 3), + image, + ) + .unwrap(); + image + }) + .collect::>(); + + (frames, num_splits_h, num_splits_w) + } + + pub fn rescale( + &self, + image: &ndarray::Array3, + rescale_factor: f32, + ) -> ndarray::Array3 { + let image = image.to_owned(); + let image = image.map(|x| *x as f32 * rescale_factor); + image + } + + fn pad_image( + &self, + image: &ndarray::Array3, + output_size: (i32, i32), + constant_values: f32, + data_format: &str, + ) -> (ndarray::Array3, ndarray::Array2) { + let (input_height, input_width, _) = image.dim(); + let (output_height, output_width) = output_size; + + // Create padded image with zeros + let mut padded_image = if data_format == "channels_first" { + ndarray::Array3::zeros((3, output_height as usize, output_width as usize)) + } else { + ndarray::Array3::zeros((output_height as usize, output_width as usize, 3)) + }; + + // Create pixel attention mask + let mut pixel_mask = + ndarray::Array2::zeros((output_height as usize, output_width as usize)); + + // Copy the original image into the padded image + if data_format == "channels_first" { + let transposed_image = image.to_owned().permuted_axes([2, 0, 1]); + padded_image + .slice_mut(s![.., ..input_height, ..input_width]) + .assign(&transposed_image); + pixel_mask + .slice_mut(s![..input_height, ..input_width]) + .fill(1); + } else { + padded_image + .slice_mut(s![..input_height, ..input_width, ..]) + .assign(&image); + pixel_mask + .slice_mut(s![..input_height, ..input_width]) + .fill(1); + } + + (padded_image, pixel_mask) + } + + pub fn pad( + &self, + images: &[Vec>], + constant_values: f32, + return_pixel_mask: bool, + data_format: &str, + ) -> (Vec>, Option>>) { + // Get max dimensions across all images + let mut max_height = 0; + let mut max_width = 0; + let mut max_num_images = 0; + + for batch in images { + max_num_images = std::cmp::max(max_num_images, batch.len()); + for image in batch { + let (height, width, _) = image.dim(); + println!("Image Shape: {:?}", image.shape()); + max_height = std::cmp::max(max_height, height); + max_width = std::cmp::max(max_width, width); + } + } + + let output_size = (max_height as i32, max_width as i32); + let batch_size = images.len(); + + // Create empty padded images and masks + let mut padded_images = vec![ + vec![ + if data_format == "channels_first" { + ndarray::Array3::zeros((3, max_height, max_width)) + } else { + ndarray::Array3::zeros((max_height, max_width, 3)) + }; + max_num_images + ]; + batch_size + ]; + + let mut padded_masks = if return_pixel_mask { + Some(vec![ + vec![ + ndarray::Array2::zeros((max_height, max_width)); + max_num_images + ]; + batch_size + ]) + } else { + None + }; + + // Pad each image + for (batch_idx, batch) in images.iter().enumerate() { + for (sample_idx, image) in batch.iter().enumerate() { + let (padded_image, pixel_mask) = + self.pad_image(image, output_size, constant_values, data_format); + padded_images[batch_idx][sample_idx] = padded_image; + if let Some(ref mut masks) = padded_masks { + masks[batch_idx][sample_idx] = pixel_mask; + } + } + } + + ( + padded_images.into_iter().flatten().collect(), + padded_masks.map(|masks| masks.into_iter().flatten().collect()), + ) + } + + fn _preprocess_one_image( + &self, + image: &DynamicImage, + ) -> Result<(ndarray::Array4, Option>, i32, i32), anyhow::Error> { + // Step 1: Initial resize + let resized_image = self.resize(image, self.size.clone().unwrap(), FilterType::Lanczos3); + + // Convert back to DynamicImage for further processing + let resized_dynamic_image = DynamicImage::ImageRgb8( + RgbImage::from_raw( + resized_image.shape()[1] as u32, + resized_image.shape()[0] as u32, + resized_image.into_raw_vec_and_offset().0, + ) + .ok_or_else(|| anyhow::anyhow!("Failed to convert resized image to DynamicImage"))?, + ); + + // Step 2: Resize for vision encoder + let vision_encoder_image = self.resize_for_vision_encoder( + &resized_dynamic_image, + &self.max_image_size.clone().unwrap()["longest_edge"], + FilterType::Lanczos3, + ); + + let vision_encoder_image = DynamicImage::ImageRgb8( + RgbImage::from_raw( + vision_encoder_image.shape()[1] as u32, + vision_encoder_image.shape()[0] as u32, + vision_encoder_image.into_raw_vec_and_offset().0, + ) + .ok_or_else(|| anyhow::anyhow!("Failed to convert resized image to DynamicImage"))?, + ); + + // Step 3: Split image if needed + let (frames, n_rows, n_cols) = if self.do_image_splitting { + self.split_image( + &vision_encoder_image, + &self.max_image_size.clone().unwrap(), + FilterType::Lanczos3, + ) + } else { + let frame = vision_encoder_image.to_rgb8().into_raw(); + let frame = ndarray::Array3::from_shape_vec( + ( + vision_encoder_image.dimensions().0 as usize, + vision_encoder_image.dimensions().1 as usize, + 3, + ), + frame, + ) + .unwrap(); + (vec![frame], 0, 0) + }; + + // Step 4: Rescale frames + let rescale_image_frames: Vec> = if self.do_rescale { + frames + .iter() + .map(|frame| self.rescale(frame, self.rescale_factor)) + .collect() + } else { + frames + .iter() + .map(|frame| frame.map(|x| *x as f32)) + .collect() + }; + + // Step 5: Normalize frames + let normalized_frames = if self.do_normalize { + let image_mean = self + .image_mean + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing image_mean"))?; + let image_std = self + .image_std + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing image_std"))?; + let image_mean = ndarray::Array3::from_shape_vec((1, 1, image_mean.len()), image_mean)?; + let image_std = ndarray::Array3::from_shape_vec((1, 1, image_std.len()), image_std)?; + + rescale_image_frames + .iter() + .map(|frame| normalize(frame, &image_mean, &image_std)) + .collect() + } else { + rescale_image_frames + }; + + // Step 6: Pad and stack frames + let (padded_images, padded_masks) = if self.do_pad { + self.pad(&[normalized_frames], 0.0, true, "channels_first") + } else { + (normalized_frames, None) + }; + + // Stack frames into a single batch + let image_views: Vec<_> = padded_images.iter().map(|arr| arr.view()).collect(); + let padded_images_concatenated = ndarray::stack(ndarray::Axis(0), &image_views)?; + + // Stack masks if they exist + let padded_masks_concatenated = if let Some(masks) = padded_masks { + let mask_views: Vec<_> = masks.iter().map(|arr| arr.view()).collect(); + Some(ndarray::stack(ndarray::Axis(0), &mask_views)?) + } else { + None + }; + + Ok(( + padded_images_concatenated, + padded_masks_concatenated, + n_rows, + n_cols, + )) + } + + pub fn preprocess( + &self, + images: &[DynamicImage], + ) -> Result<(ndarray::Array4, Option>, i32, i32), anyhow::Error> { + let mut preprocessed_images = Vec::new(); + let mut preprocessed_masks = Vec::new(); + let mut n_rows = 0; + let mut n_cols = 0; + for image in images { + let (padded_images, padded_masks, rows, cols) = self._preprocess_one_image(image)?; + preprocessed_images.push(padded_images); + if let Some(mask) = padded_masks { + preprocessed_masks.push(mask); + } + n_rows = rows; + n_cols = cols; + } + let image_views: Vec<_> = preprocessed_images.iter().map(|arr| arr.view()).collect(); + let preprocessed_images = ndarray::concatenate(ndarray::Axis(0), &image_views)?; + let preprocessed_masks = if !preprocessed_masks.is_empty() { + let mask_views: Vec<_> = preprocessed_masks.iter().map(|arr| arr.view()).collect(); + Some(ndarray::concatenate(ndarray::Axis(0), &mask_views)?) + } else { + None + }; + Ok((preprocessed_images, preprocessed_masks, n_rows, n_cols)) + } +} + +pub struct Idefics3Processor { + image_processor: Idefics3ImageProcessor, + tokenizer: Tokenizer, + regex: Regex, + fake_image_token: AddedToken, + image_token: AddedToken, + end_of_utterance_token: AddedToken, + global_image_tag: AddedToken, +} + +impl Idefics3Processor { + pub fn from_pretrained(model_id: &str) -> anyhow::Result { + let image_processor = Idefics3ImageProcessor::from_pretrained(model_id)?; + let mut tokenizer = Tokenizer::from_pretrained(model_id, None) + .map_err(|e| anyhow::anyhow!("Tokenizer error: {}", e))?; + let fake_image_token = AddedToken::from("", true); + let image_token = AddedToken::from("", true); + let end_of_utterance_token = AddedToken::from("", true); + let global_image_tag = AddedToken::from("", true); + tokenizer.add_special_tokens(&[ + fake_image_token.clone(), + image_token.clone(), + end_of_utterance_token.clone(), + global_image_tag.clone(), + ]); + + let regex = Regex::new(r"(\n?\n?|\n?)+").unwrap(); + Ok(Idefics3Processor { + image_processor, + tokenizer, + regex, + fake_image_token, + image_token, + end_of_utterance_token, + global_image_tag, + }) + } + + pub fn preprocess( + &self, + images: &[DynamicImage], + ) -> anyhow::Result<( + ndarray::Array2, + ndarray::Array2, + ndarray::Array4, + Option>, + )> { + let (preprocessed_images, preprocessed_masks, n_rows, n_cols) = + self.image_processor.preprocess(images)?; + println!("N rows: {:?}, N cols: {:?}", n_rows, n_cols); + let image_prompt = _prompt_split_image( + 169, + n_rows, + n_cols, + &self.fake_image_token, + &self.image_token, + &self.global_image_tag, + ); + + let prompt = "<|im_start|>user\nDescribe the image."; + + // in the prompt replace the image_token with the image_prompt + let prompt = prompt.replace(&self.image_token.content, &image_prompt); + + let encodings = self + .tokenizer + .encode(prompt, true) + .map_err(|e| anyhow::anyhow!("Tokenizer error: {}", e))?; + + let input_ids = Array2::from_shape_vec( + (1, encodings.get_ids().len()), + encodings.get_ids().iter().map(|x| *x as i64).collect(), + )?; + let attention_mask = Array2::from_shape_vec( + (1, encodings.get_attention_mask().len()), + encodings + .get_attention_mask() + .iter() + .map(|x| *x as i64) + .collect(), + )?; + + // Count occurrences of token ID 49190 in the first row + let token_count = input_ids.row(0).iter().filter(|&&x| x == 49190).count(); + + Ok(( + input_ids, + attention_mask, + preprocessed_images, + preprocessed_masks, + )) + } +} + +fn _prompt_split_image( + image_seq_len: i32, + image_rows: i32, + image_cols: i32, + fake_token_around_image: &AddedToken, + image_token: &AddedToken, + global_img_token: &AddedToken, +) -> String { + let mut text_split_images = String::new(); + for n_h in 0..image_rows { + for n_w in 0..image_cols { + text_split_images.push_str(&format!( + "{}{}", + fake_token_around_image.content, + n_h + 1, + n_w + 1, + image_token.content.repeat(image_seq_len as usize), + )); + } + text_split_images.push_str("\n"); + } + text_split_images.push_str(&format!( + "\n{}{}{}{}", + fake_token_around_image.content, + global_img_token.content, + image_token.content.repeat(image_seq_len as usize), + fake_token_around_image.content + )); + text_split_images +} + +fn get_resize_output_image_size(image: DynamicImage, resolution_max_side: i32) -> (i32, i32) { + let (width, height) = image.dimensions(); + let (new_height, new_width) = + _resize_output_size_rescale_to_max_len(height as i32, width as i32, 1, resolution_max_side); + let (new_height, new_width) = + _resize_output_size_scale_below_upper_bound(new_height, new_width, Some(MAX_IMAGE_SIZE)); + (new_height as i32, new_width as i32) +} + +fn _resize_output_size_rescale_to_max_len( + height: i32, + width: i32, + min_len: i32, + max_len: i32, +) -> (i32, i32) { + let max_len = if max_len == 0 { + std::cmp::max(height, width) + } else { + max_len + }; + let aspect_ratio = width as f32 / height as f32; + let (new_width, new_height) = if width >= height { + let new_width = max_len; + let new_height = (new_width as f32 / aspect_ratio) as i32; + if new_height % 2 != 0 { + (new_width, new_height + 1) + } else { + (new_width, new_height) + } + } else { + let new_height = max_len; + let new_width = (new_height as f32 * aspect_ratio) as i32; + if new_width % 2 != 0 { + (new_width + 1, new_height) + } else { + (new_width, new_height) + } + }; + + // Avoid resizing to a size smaller than min_len + let new_height = std::cmp::max(new_height, min_len); + let new_width = std::cmp::max(new_width, min_len); + (new_height, new_width) +} + +fn _resize_output_size_scale_below_upper_bound( + height: i32, + width: i32, + max_len: Option, +) -> (i32, i32) { + let max_len = max_len.unwrap_or_else(|| std::cmp::max(height, width)); + let aspect_ratio = width as f32 / height as f32; + + let (new_width, new_height) = if width >= height && width > max_len { + let new_width = max_len; + let new_height = (new_width as f32 / aspect_ratio) as i32; + (new_width, new_height) + } else if height > width && height > max_len { + let new_height = max_len; + let new_width = (new_height as f32 * aspect_ratio) as i32; + (new_width, new_height) + } else { + (width, height) + }; + + // Avoid resizing to a size smaller than 1 + let new_height = std::cmp::max(new_height, 1); + let new_width = std::cmp::max(new_width, 1); + (new_height, new_width) +} + +fn normalize( + image: &ndarray::Array3, + mean: &ndarray::Array3, + std: &ndarray::Array3, +) -> ndarray::Array3 { + let normalized_image = (image - mean) / std; + + normalized_image +} + +#[cfg(test)] +mod tests { + use super::*; + use hf_hub::api::sync::Api; + use hf_hub::{Repo, RepoType}; + use image::RgbImage; + + #[test] + fn image_resize_test() { + let image = image::open("/home/akshay/projects/EmbedAnything/test.jpg").unwrap(); + let image_array = image.to_rgb8().into_raw(); + + let api = Api::new().unwrap(); + let repo = api.repo(Repo::new( + "onnx-community/colSmol-256M-ONNX".to_string(), + RepoType::Model, + )); + let config_file = repo.get("preprocessor_config.json").unwrap(); + let processor: Idefics3ImageProcessor = + serde_json::from_slice(&std::fs::read(config_file).unwrap()).unwrap(); + println!("{:?}", processor); + let resized_image = processor.resize( + &image, + processor.size.clone().unwrap(), + FilterType::Lanczos3, + ); + // println!("Resized Image: {:?}", resized_image.into_raw_vec_and_offset().0.len()); + // println!("Resized Image: {:?}", resized_image); + + let resized_dynamic_image = DynamicImage::ImageRgb8( + RgbImage::from_raw( + resized_image.shape()[1] as u32, + resized_image.shape()[0] as u32, + resized_image.into_raw_vec_and_offset().0, + ) + .unwrap(), + ); + println!( + "Resized Dynamic Image: {:?}", + resized_dynamic_image.dimensions() + ); + let resized_image = processor.resize_for_vision_encoder( + &resized_dynamic_image, + &processor.max_image_size.clone().unwrap()["longest_edge"], + FilterType::Lanczos3, + ); + println!("Resized Image: {:?}", resized_image.shape()); + + let (frames, num_splits_h, num_splits_w) = processor.split_image( + &resized_dynamic_image, + &processor.max_image_size.clone().unwrap(), + FilterType::Lanczos3, + ); + println!("Frames: {:?}", frames.len()); + println!( + "Num Splits H: {:?}, Num Splits W: {:?}", + num_splits_h, num_splits_w + ); + + let rescale_image_frames = frames + .iter() + .map(|frame| { + let frame = frame.to_owned(); + processor.rescale(&frame, processor.rescale_factor) + }) + .collect::>(); + // println!("Rescale Image: {:?}", rescale_image); + + let image_mean = processor.image_mean.clone().unwrap(); + let image_std = processor.image_std.clone().unwrap(); + let image_mean = + ndarray::Array3::from_shape_vec((1, 1, image_mean.len()), image_mean).unwrap(); + let image_std = + ndarray::Array3::from_shape_vec((1, 1, image_std.len()), image_std).unwrap(); + let normalized_frames = rescale_image_frames + .iter() + .map(|frame| { + let frame = frame.to_owned(); + normalize(&frame, &image_mean, &image_std) + }) + .collect::>(); + + // println!("Normalized Frames: {:?}", normalized_frames.len()); + let (padded_images, padded_masks) = + processor.pad(&[normalized_frames], 0.0, true, "channels_first"); + let image_views: Vec<_> = padded_images.iter().map(|arr| arr.view()).collect(); + let padded_images_concatenated = ndarray::stack(ndarray::Axis(0), &image_views).unwrap(); + println!( + "Padded Images Concatenated: {:?}", + padded_images_concatenated.shape() + ); + if let Some(masks) = padded_masks { + let mask_views: Vec<_> = masks.iter().map(|arr| arr.view()).collect(); + let padded_masks_concatenated = ndarray::stack(ndarray::Axis(0), &mask_views); + println!("Padded Masks Concatenated: {:?}", padded_masks_concatenated); + } + } + + #[test] + fn test_idefics3_processor() { + let image = image::open("/home/akshay/projects/EmbedAnything/test.jpg").unwrap(); + let processor = + Idefics3Processor::from_pretrained("onnx-community/colSmol-256M-ONNX").unwrap(); + processor.preprocess(&[image]).unwrap(); + } +} diff --git a/rust/src/models/idefics3/mod.rs b/rust/src/models/idefics3/mod.rs new file mode 100644 index 00000000..0f38153f --- /dev/null +++ b/rust/src/models/idefics3/mod.rs @@ -0,0 +1,2 @@ +pub mod model; +pub mod tensor_processing; diff --git a/rust/src/models/idefics3/model.rs b/rust/src/models/idefics3/model.rs new file mode 100644 index 00000000..e9caf4e9 --- /dev/null +++ b/rust/src/models/idefics3/model.rs @@ -0,0 +1,820 @@ +use crate::models::llama::{self, Cache, LlamaBase}; +use crate::models::with_tracing::{linear, linear_no_bias, Embedding, Linear}; +use candle_core::{CpuStorage, CustomOp1, DType, Device, Layout, Module, Shape, WithDType, D}; +use candle_core::{Result, Tensor}; +use candle_nn::{Conv2dConfig, LayerNorm, LayerNormConfig}; +use serde::{Deserialize, Serialize}; + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} +struct NonZero {} + +impl NonZero { + // Sequential version + fn nonzero(&self, vs: &[T], layout: &Layout) -> Vec { + let n = layout.dims().len(); + let mut result = Vec::new(); + let mut indices = vec![0u32; n]; + for (i, v) in vs.iter().enumerate() { + if !v.is_zero() { + let mut idx = i; + for (dim_index, dim) in layout.dims().iter().enumerate().rev() { + let d = idx % dim; + indices[dim_index] = u32::try_from(d).unwrap(); + idx /= dim; + } + result.extend_from_slice(&indices); + } + } + result + } +} + +impl CustomOp1 for NonZero { + fn name(&self) -> &'static str { + "nonzero" + } + + fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)> { + if !layout.is_contiguous() { + return Err(candle_core::Error::RequiresContiguous { op: "nonzero" }); + } + let result = match storage { + CpuStorage::U8(vs) => self.nonzero(vs, layout), + CpuStorage::U32(vs) => self.nonzero(vs, layout), + CpuStorage::I64(vs) => self.nonzero(vs, layout), + CpuStorage::BF16(vs) => self.nonzero(vs, layout), + CpuStorage::F16(vs) => self.nonzero(vs, layout), + CpuStorage::F32(vs) => self.nonzero(vs, layout), + CpuStorage::F64(vs) => self.nonzero(vs, layout), + }; + let index_len = layout.dims().len(); + let result_len = result.len() / index_len; + let result = CpuStorage::U32(result); + let shape = Shape::from_dims(&[result_len, index_len]); + Ok((result, shape)) + } +} + +pub trait NonZeroOp { + fn nonzero(&self) -> Result; +} + +impl NonZeroOp for Tensor { + fn nonzero(&self) -> Result { + if !self.is_contiguous() { + return Err(candle_core::Error::RequiresContiguous { op: "nonzero" }); + } + let original_device = self.device(); + self.to_device(&candle_core::Device::Cpu)? + .apply_op1_no_bwd(&NonZero {})? + .to_device(original_device) + } +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Idefic3VisionConfig { + hidden_size: usize, + intermediate_size: usize, + num_hidden_layers: usize, + num_attention_heads: usize, + num_channels: usize, + pub patch_size: usize, + image_size: usize, + attention_dropout: f32, + layer_norm_eps: f64, + hidden_act: candle_nn::Activation, + initializer_range: f64, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct Idefics3Config { + pub vision_config: Idefic3VisionConfig, + pub text_config: llama::LlamaConfig, + pub scale_factor: Option, + pub image_token_id: usize, +} + +struct Idefics3VisionEmbeddings { + patch_size: usize, + patch_embeddings: candle_nn::Conv2d, + num_patches_per_side: usize, + position_embeddings: Embedding, +} + +impl Idefics3VisionEmbeddings { + fn load(config: &Idefic3VisionConfig, vb: candle_nn::VarBuilder) -> Result { + let embed_dim = config.hidden_size; + let image_size = config.image_size; + let patch_size = config.patch_size; + let num_patches_per_side = image_size / patch_size; + let num_patches = num_patches_per_side * num_patches_per_side; + let num_position = num_patches; + let patch_embeddings = candle_nn::conv2d( + config.num_channels, + config.hidden_size, + config.patch_size, + Conv2dConfig { + stride: config.patch_size, + padding: 0, + groups: 1, + dilation: 1, + cudnn_fwd_algo: None, + }, + vb.pp("patch_embedding"), + )?; + let position_embeddings = + Embedding::new(num_position, embed_dim, vb.pp("position_embedding"))?; + Ok(Self { + patch_size, + patch_embeddings, + num_patches_per_side, + position_embeddings, + }) + } + + fn forward( + &self, + pixel_values: &Tensor, + patch_attention_mask: &Tensor, + device: &Device, + ) -> Result { + let batch_size = pixel_values.dims()[0]; + let max_im_h = pixel_values.dims()[2]; + let max_im_w = pixel_values.dims()[3]; + + let patch_embeds = self.patch_embeddings.forward(pixel_values)?; + let embeddings = patch_embeds.flatten_from(2)?.transpose(1, 2)?; + let (max_nb_patchs_h, max_nb_patchs_w) = + (max_im_h / self.patch_size, max_im_w / self.patch_size); + let boundaries = Tensor::arange_step( + 1.0 / self.num_patches_per_side as f64, + 1.0, + 1.0 / self.num_patches_per_side as f64, + &Device::Cpu, + )? + .to_vec1::()?; + let mut position_ids = Tensor::zeros( + (batch_size, max_nb_patchs_h * max_nb_patchs_w), + DType::I64, + device, + )?; + + for batch_idx in 0..batch_size { + let p_attn_mask = patch_attention_mask.get(batch_idx)?; + let nb_patches_h = p_attn_mask.get_on_dim(1, 0)?.sum_all()?.to_scalar::()?; + let nb_patches_w = p_attn_mask.get_on_dim(1, 1)?.sum_all()?.to_scalar::()?; + + let fractional_coords_h = + Tensor::arange_step(0., 1. - 1e-6, 1. / nb_patches_h as f64, device)? + .to_vec1::()?; + let fractional_coords_w = + Tensor::arange_step(0., 1. - 1e-6, 1. / nb_patches_w as f64, device)? + .to_vec1::()?; + + let bucket_coords_h = bucketize(&fractional_coords_h, &boundaries, true); + let bucket_coords_w = bucketize(&fractional_coords_w, &boundaries, true); + + let bucket_coords_h_tensor = + Tensor::from_vec(bucket_coords_h.clone(), (bucket_coords_h.len(),), device)?; + let bucket_coords_w_tensor = + Tensor::from_vec(bucket_coords_w.clone(), (bucket_coords_w.len(),), device)?; + + let pos_ids = (bucket_coords_h_tensor.unsqueeze(1)? + * (self.num_patches_per_side as f64))? + .broadcast_add(&bucket_coords_w_tensor)? + .flatten_from(0)?; + + let p_attn_mask_flat = p_attn_mask.flatten_from(0)?; + // Use tensor operations to find indices where mask is 1 + let indices = p_attn_mask_flat + .to_dtype(DType::F32)? + .nonzero()? + .squeeze(1)?; + + let positions = pos_ids.gather(&indices, 0)?; + position_ids = position_ids.slice_assign( + &[batch_idx..batch_idx + 1, 0..positions.dims()[0]], + &positions.unsqueeze(0)?, + )?; + } + let position_embeddings = self.position_embeddings.forward(&position_ids)?; + let embeddings = embeddings.add(&position_embeddings)?; + + Ok(embeddings) + } +} + +struct Idefics3VisionAttention { + num_heads: usize, + head_dim: usize, + scale: f32, + k_proj: Linear, + v_proj: Linear, + q_proj: Linear, + out_proj: Linear, + is_causal: bool, + use_flash_attn: bool, +} + +impl Idefics3VisionAttention { + fn load( + config: &Idefic3VisionConfig, + use_flash_attn: bool, + vb: candle_nn::VarBuilder, + ) -> Result { + let embed_dim = config.hidden_size; + let num_heads = config.num_attention_heads; + let head_dim = embed_dim / num_heads; + let scale = (head_dim as f32).powf(-0.5); + let k_proj = linear(embed_dim, embed_dim, vb.pp("k_proj"))?; + let v_proj = linear(embed_dim, embed_dim, vb.pp("v_proj"))?; + let q_proj = linear(embed_dim, embed_dim, vb.pp("q_proj"))?; + let out_proj = linear(embed_dim, embed_dim, vb.pp("out_proj"))?; + Ok(Self { + num_heads, + head_dim, + scale, + k_proj, + v_proj, + q_proj, + out_proj, + is_causal: false, + use_flash_attn, + }) + } + + fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Option, + ) -> Result<(Tensor, Option)> { + let (batch_size, q_len, embed_dim) = hidden_states.dims3()?; + let q = self.q_proj.forward(hidden_states)?; + let k = self.k_proj.forward(hidden_states)?; + let v = self.v_proj.forward(hidden_states)?; + + if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let scale = 1f32 / (self.head_dim as f32).sqrt(); + let attn_output = flash_attn(&q, &k, &v, scale, self.is_causal)?.transpose(1, 2)?; + let attn_output = attn_output + .transpose(1, 2)? + .reshape((batch_size, q_len, embed_dim))?; + let attn_output = self.out_proj.forward(&attn_output)?; + return Ok((attn_output, None)); + } + let shape = (batch_size, q_len, self.num_heads, self.head_dim); + + let query_states = q.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let key_states = k.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let value_states = v.reshape(shape)?.transpose(1, 2)?.contiguous()?; + + let attn_weights = (query_states.matmul(&key_states.t()?)? * self.scale as f64)?; + let attn_weights = match attention_mask { + None => attn_weights, + Some(mask) => attn_weights.broadcast_add(mask)?, + }; + // The original implementation upcasts to f32 but candle_nn::ops::softmax should handle this properly. + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_outputs = attn_weights + .matmul(&value_states)? + .transpose(1, 2)? + .reshape((batch_size, q_len, ()))? + .apply(&self.out_proj)?; + Ok((attn_outputs, None)) + } +} + +struct Idefics3VisionMLP { + activation_fn: candle_nn::Activation, + fc1: Linear, + fc2: Linear, +} + +impl Idefics3VisionMLP { + fn load(config: &Idefic3VisionConfig, vb: candle_nn::VarBuilder) -> Result { + let activation_fn = config.hidden_act; + let fc1 = linear(config.hidden_size, config.intermediate_size, vb.pp("fc1"))?; + let fc2 = linear(config.intermediate_size, config.hidden_size, vb.pp("fc2"))?; + Ok(Self { + activation_fn, + fc1, + fc2, + }) + } + + fn forward(&self, hidden_states: &Tensor) -> Result { + let hidden_states = self.fc1.forward(hidden_states)?; + let hidden_states = self.activation_fn.forward(&hidden_states)?; + let hidden_states = self.fc2.forward(&hidden_states)?; + Ok(hidden_states) + } +} + +struct Idefics3SimpleMLP { + proj: Linear, +} + +impl Idefics3SimpleMLP { + fn load(config: &Idefics3Config, vb: candle_nn::VarBuilder) -> Result { + let proj = linear_no_bias( + config.vision_config.hidden_size * (config.scale_factor.unwrap_or(2).pow(2)), + config.text_config.hidden_size, + vb.pp("proj"), + )?; + Ok(Self { proj }) + } + + fn forward(&self, hidden_states: &Tensor) -> Result { + let hidden_states = self.proj.forward(hidden_states)?; + Ok(hidden_states) + } +} + +struct Idefics3EncoderLayer { + self_attn: Idefics3VisionAttention, + layer_norm1: candle_nn::LayerNorm, + layer_norm2: candle_nn::LayerNorm, + mlp: Idefics3VisionMLP, +} + +impl Idefics3EncoderLayer { + fn load( + config: &Idefic3VisionConfig, + use_flash_attn: bool, + vb: candle_nn::VarBuilder, + ) -> Result { + let self_attn = Idefics3VisionAttention::load(config, use_flash_attn, vb.pp("self_attn"))?; + let layer_norm1 = candle_nn::layer_norm( + config.hidden_size, + LayerNormConfig::from(config.layer_norm_eps), + vb.pp("layer_norm1"), + )?; + let layer_norm2 = candle_nn::layer_norm( + config.hidden_size, + LayerNormConfig::from(config.layer_norm_eps), + vb.pp("layer_norm2"), + )?; + let mlp = Idefics3VisionMLP::load(config, vb.pp("mlp"))?; + Ok(Self { + self_attn, + layer_norm1, + layer_norm2, + mlp, + }) + } + + pub fn forward( + &self, + hidden_states: &Tensor, + attention_mask: &Option, + ) -> Result<(Tensor, Option)> { + let residual = hidden_states; + + let hidden_states = hidden_states.apply(&self.layer_norm1)?; + let (hidden_states, attn_weights) = + self.self_attn.forward(&hidden_states, attention_mask)?; + let hidden_states = hidden_states.add(residual)?; + + let residual = hidden_states.clone(); + let hidden_states = hidden_states.apply(&self.layer_norm2)?; + let hidden_states = self.mlp.forward(&hidden_states)?.add(&residual)?; + Ok((hidden_states, attn_weights)) + } +} + +struct Idefics3Encoder { + layers: Vec, +} + +impl Idefics3Encoder { + fn load( + config: &Idefic3VisionConfig, + use_flash_attn: bool, + vb: candle_nn::VarBuilder, + ) -> Result { + let layers = (0..config.num_hidden_layers) + .map(|i| { + Idefics3EncoderLayer::load(config, use_flash_attn, vb.pp(format!("layers.{}", i))) + }) + .collect::>>()?; + Ok(Self { layers }) + } + + fn forward(&self, input_embeds: &Tensor, attention_mask: &Option) -> Result { + let mut hidden_states = input_embeds.clone(); + for layer in &self.layers { + hidden_states = layer.forward(&hidden_states, attention_mask)?.0; + } + Ok(hidden_states) + } +} + +struct Idefics3Connector { + modaliity_projection: Idefics3SimpleMLP, + scale_factor: Option, +} + +impl Idefics3Connector { + pub fn load(config: &Idefics3Config, vb: candle_nn::VarBuilder) -> Result { + let modaliity_projection = Idefics3SimpleMLP::load(config, vb.pp("modality_projection"))?; + Ok(Self { + modaliity_projection, + scale_factor: config.scale_factor, + }) + } + + pub fn pixel_shuffle(&self, x: &Tensor, scale_factor: Option) -> Result { + let scale_factor = scale_factor.unwrap_or(2); + let (b_sz, seq, embed_dim) = x.dims3()?; + let height = (seq as f64).sqrt() as usize; + let width = height; + + let x = x.reshape((b_sz, height, width, embed_dim))?; + let x = x.reshape((b_sz, height, width / scale_factor, embed_dim * scale_factor))?; + let x = x.permute((0, 2, 1, 3))?; + let x = x.reshape(( + b_sz, + width / scale_factor, + height / scale_factor, + embed_dim * scale_factor * scale_factor, + ))?; + let x = x.permute((0, 2, 1, 3))?; + let x = x.reshape(( + b_sz, + seq / (scale_factor * scale_factor), + embed_dim * scale_factor * scale_factor, + ))?; + Ok(x) + } + + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.pixel_shuffle(x, self.scale_factor)?; + let x = self.modaliity_projection.forward(&x)?; + Ok(x) + } +} + +struct Idefics3VisionTransformer { + embeddings: Idefics3VisionEmbeddings, + encoder: Idefics3Encoder, + post_layernorm: LayerNorm, +} + +impl Idefics3VisionTransformer { + fn load( + config: &Idefics3Config, + use_flash_attn: bool, + vb: candle_nn::VarBuilder, + ) -> Result { + let embeddings = + Idefics3VisionEmbeddings::load(&config.vision_config, vb.pp("embeddings"))?; + let encoder = + Idefics3Encoder::load(&config.vision_config, use_flash_attn, vb.pp("encoder"))?; + let post_layernorm = candle_nn::layer_norm( + config.vision_config.hidden_size, + LayerNormConfig::from(config.vision_config.layer_norm_eps), + vb.pp("post_layernorm"), + )?; + Ok(Self { + embeddings, + encoder, + post_layernorm, + }) + } + + fn forward( + &self, + pixel_values: &Tensor, + patch_attention_mask: Option<&Tensor>, + ) -> Result { + let patch_attention_mask = if let Some(patch_attention_mask) = patch_attention_mask { + patch_attention_mask.clone() + } else { + Tensor::ones( + ( + pixel_values.dims()[0], + pixel_values.dims()[2], + pixel_values.dims()[3], + ), + DType::F32, + pixel_values.device(), + )? + }; + let hidden_states = + self.embeddings + .forward(pixel_values, &patch_attention_mask, pixel_values.device())?; + let patch_attention_mask = patch_attention_mask.flatten_from(1)?; + let patch_attention_mask = + prepare_4d_attention_mask(&patch_attention_mask, pixel_values.dtype(), None)?; + + let hidden_states = self + .encoder + .forward(&hidden_states, &Some(patch_attention_mask))?; + let hidden_states = self.post_layernorm.forward(&hidden_states)?; + Ok(hidden_states) + } +} + +pub struct Idefics3Model { + vision_model: Idefics3VisionTransformer, + connector: Idefics3Connector, + pub text_model: LlamaBase, + image_token_id: usize, + use_flash_attn: bool, + config: Idefics3Config, + dtype: DType, +} + +impl Idefics3Model { + pub fn load( + config: &Idefics3Config, + use_flash_attn: bool, + vb: candle_nn::VarBuilder, + ) -> Result { + let vision_model = + Idefics3VisionTransformer::load(config, use_flash_attn, vb.pp("vision_model"))?; + let connector = Idefics3Connector::load(config, vb.pp("connector"))?; + let text_model = LlamaBase::load( + vb.pp("text_model"), + &config.text_config.clone().into_config(use_flash_attn), + )?; + + let image_token_id = config.image_token_id; + Ok(Self { + vision_model, + connector, + text_model, + image_token_id, + use_flash_attn, + config: config.clone(), + dtype: vb.dtype(), + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + pixel_values: &Option, + pixel_attention_mask: &Option, + kv_cache: &mut Cache, + ) -> Result { + if let (Some(pixel_values), Some(pixel_attention_mask)) = + (pixel_values, pixel_attention_mask) + { + let pixel_values = pixel_values.to_dtype(self.dtype)?; + + let input_embeds = self.text_model.embed(input_ids)?; + + let (bsz, num_images, channels, height, width) = pixel_values.dims5()?; + let pixel_values = pixel_values.reshape((bsz * num_images, channels, height, width))?; + let pixel_attention_mask = + pixel_attention_mask.reshape((bsz * num_images, height, width))?; + + let nb_values_per_image = + pixel_values.dims()[1] * pixel_values.dims()[2] * pixel_values.dims()[3]; + let real_image_inds = pixel_values + .eq(0.0)? + .sum((3, 2, 1))? + .ne(nb_values_per_image as f64)?; + + let indices = real_image_inds + .to_dtype(DType::F32)? + .eq(1.0)? + .nonzero()? + .squeeze(1)?; + + let pixel_values = pixel_values.index_select(&indices, 0)?; + let pixel_attention_mask = pixel_attention_mask.index_select(&indices, 0)?; + + let patches_subgrid = unfold( + &pixel_attention_mask, + self.config.vision_config.patch_size, + self.config.vision_config.patch_size, + 1, + ); + + let patches_subgrid = unfold( + &patches_subgrid?, + self.config.vision_config.patch_size, + self.config.vision_config.patch_size, + 2, + )?; + + let patch_attention_mask = patches_subgrid.sum((D::Minus1, D::Minus2))?.ge(0.0)?; + + let image_hidden_states = self + .vision_model + .forward(&pixel_values, Some(&patch_attention_mask))?; + + let image_hidden_states = self.connector.forward(&image_hidden_states)?; + + let new_input_embeds = + self.inputs_merger(input_ids, &input_embeds, &image_hidden_states)?; + + let output = self + .text_model + .forward_input_embed(&new_input_embeds, 0, kv_cache)?; + + Ok(output) + } else { + self.text_model.forward(input_ids, 0, kv_cache) + } + } + + fn inputs_merger( + &self, + input_ids: &Tensor, + input_embeds: &Tensor, + image_hidden_states: &Tensor, + ) -> Result { + let (num_images, seq_len, vision_hidden_size) = image_hidden_states.dims3()?; + let (bsz, text_seq_len, embed_dim) = input_embeds.dims3()?; + + let input_embeds_reshaped = input_embeds.reshape((bsz * text_seq_len, embed_dim))?; + + let input_ids = input_ids.flatten_from(0)?; + + let image_hidden_states = + image_hidden_states.reshape((num_images * seq_len, vision_hidden_size))?; + + let special_image_token_indices = input_ids + .eq(self.image_token_id as f64)? + .nonzero()? + .repeat((1, embed_dim))?; + + let new_input_embeds = + input_embeds_reshaped.scatter(&special_image_token_indices, &image_hidden_states, 0)?; + let new_input_embeds = new_input_embeds.reshape((bsz, text_seq_len, embed_dim))?; + + Ok(new_input_embeds) + } +} + +pub struct ColIdefics3Model { + model: Idefics3Model, + linear: Linear, + dtype: DType, +} + +impl ColIdefics3Model { + pub fn load( + config: &Idefics3Config, + use_flash_attn: bool, + vb: candle_nn::VarBuilder, + ) -> Result { + let model = Idefics3Model::load(config, use_flash_attn, vb.pp("model"))?; + let dim = 128; + let linear = linear(model.config.text_config.hidden_size, dim, vb.pp("linear"))?; + Ok(Self { + model, + linear, + dtype: vb.dtype(), + }) + } + + pub fn forward( + &self, + input_ids: &Tensor, + attention_mask: &Tensor, + pixel_values: &Option, + pixel_attention_mask: &Option, + ) -> Result { + let output = self.model.forward( + input_ids, + pixel_values, + pixel_attention_mask, + &mut Cache::new( + false, + self.dtype, + &self + .model + .config + .text_config + .clone() + .into_config(self.model.use_flash_attn), + input_ids.device(), + )?, + )?; + let proj = self.linear.forward(&output)?; + let proj = proj.broadcast_div(&proj.sqr()?.sum_keepdim(2)?.sqrt()?)?; + let proj = proj.broadcast_mul(&attention_mask.unsqueeze(2)?.to_dtype(proj.dtype())?)?; + + Ok(proj) + } +} + +fn bucketize(inputs: &[f64], boundaries: &[f64], right: bool) -> Vec { + // Pre-allocate with capacity for better performance + let mut result = Vec::with_capacity(inputs.len()); + + // Use binary search to find the bucket for each input + // This is O(log n) instead of O(n) for each input + for &input in inputs { + let bucket = match boundaries.binary_search_by(|&boundary| { + if input < boundary || (!right && input == boundary) { + std::cmp::Ordering::Greater + } else { + std::cmp::Ordering::Less + } + }) { + Ok(pos) => pos, + Err(pos) => pos, + }; + result.push(bucket as i64); + } + + result +} + +fn unfold(x: &Tensor, size: usize, step: usize, dim: usize) -> Result { + let x_shape = x.shape().dims().to_vec(); + let len = x_shape[dim]; + let num = (len - size) / step + 1; + + let mut idx_data = Vec::with_capacity(num * size); + for i in 0..num { + let base = i * step; + for j in 0..size { + idx_data.push((base + j) as i64); + } + } + + let mut perm: Vec = (0..x_shape.len()).filter(|&i| i != dim).collect(); + perm.push(dim); + let x = x.permute(perm)?; + + let mut inv_perm: Vec = (0..x_shape.len()).collect(); + let moved_element = inv_perm.remove(0); + inv_perm.insert(dim, moved_element); + inv_perm.push(x_shape.len()); + + let mut idx_shape = vec![num]; + for (i, _) in x_shape.iter().enumerate() { + if i != dim { + idx_shape.push(x_shape[i]); + } + } + idx_shape.push(size); + + let idx = Tensor::from_vec(idx_data, &[num, size], x.device())?; + + let mut reshape_dims = vec![num]; + for i in 0..x_shape.len() { + if i != dim { + reshape_dims.push(1); + } + } + reshape_dims.push(size); + + let reshape_dims: &[usize] = &reshape_dims; + let idx = idx + .reshape(reshape_dims)? + .broadcast_as(&idx_shape[..])? + .contiguous()?; + + let mut repeat_dims = vec![1; x_shape.len()]; + repeat_dims[0] = num; + let x = x.unsqueeze(0)?.repeat(repeat_dims)?; + + let x = x.gather(&idx, x_shape.len())?.permute(inv_perm)?; + + Ok(x) +} + +// Global attention mask calculated from padded token inputs +fn prepare_4d_attention_mask( + mask: &Tensor, + dtype: DType, + tgt_len: Option, +) -> Result { + let bsz = mask.dim(0)?; + let src_len = mask.dim(1)?; + let tgt_len = tgt_len.unwrap_or(src_len); + + let expanded_mask = mask + .unsqueeze(1)? + .unsqueeze(2)? + .expand((bsz, 1, tgt_len, src_len))?; + + let inverted_mask = (1.0 - expanded_mask)?; + + (inverted_mask * f32::MIN as f64)?.to_dtype(dtype) +} diff --git a/rust/src/models/idefics3/tensor_processing.rs b/rust/src/models/idefics3/tensor_processing.rs new file mode 100644 index 00000000..df200c3e --- /dev/null +++ b/rust/src/models/idefics3/tensor_processing.rs @@ -0,0 +1,768 @@ +use candle_core::{DType, Device, Tensor}; +use hf_hub::{api::sync::Api, Repo, RepoType}; +use image::{imageops::FilterType, DynamicImage, GenericImageView, RgbImage}; +use serde::Deserialize; +use std::collections::HashMap; +use tokenizers::{AddedToken, Tokenizer}; + +const MAX_IMAGE_SIZE: i32 = 4096; + +type ProcessedImageResult = Result<(Vec, Option>, i32, i32), anyhow::Error>; + +#[derive(Debug, Clone, Deserialize)] +pub struct Idefics3ImageProcessor { + size: Option>, + do_image_splitting: bool, + image_mean: Option>, + image_std: Option>, + + #[serde(default = "default_max_image_size")] + max_image_size: Option>, + do_rescale: bool, + rescale_factor: f32, + do_pad: bool, + do_normalize: bool, +} + +fn default_max_image_size() -> Option> { + Some(HashMap::from([("longest_edge".to_string(), 364)])) +} + +impl Idefics3ImageProcessor { + pub fn resize_for_vision_encoder( + &self, + image: &DynamicImage, + vision_encoder_max_size: &i32, + resample: FilterType, + device: &Device, + ) -> Result { + let (width, height) = image.dimensions(); + let height = height as f32; + let width = width as f32; + let vision_encoder_max_size = *vision_encoder_max_size as f32; + + let aspect_ratio = width / height; + let (new_height, new_width) = if width >= height { + let width = (width / vision_encoder_max_size).ceil() * vision_encoder_max_size; + let height = (width / aspect_ratio).floor(); + let height = + ((height / vision_encoder_max_size).ceil() * vision_encoder_max_size).ceil(); + (height, width) + } else { + let height = (height / vision_encoder_max_size).ceil() * vision_encoder_max_size; + let width = (height * aspect_ratio).floor(); + let width = (width / vision_encoder_max_size).ceil() * vision_encoder_max_size; + (height, width) + }; + self.resize( + image, + HashMap::from([ + ("width".to_string(), new_width as i32), + ("height".to_string(), new_height as i32), + ]), + resample, + device, + ) + } + + pub fn resize( + &self, + image: &DynamicImage, + size: HashMap, + resample: FilterType, + device: &Device, + ) -> Result { + let resized_image = if size.contains_key("height") && size.contains_key("width") { + image.resize_exact( + size.get("width") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing width"))? as u32, + size.get("height") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing height"))? as u32, + resample, + ) + } else { + let size = get_resize_output_image_size( + image.clone(), + size.get("longest_edge") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing longest_edge"))?, + ); + image.resize_exact(size.1 as u32, size.0 as u32, resample) + }; + let (width, height) = resized_image.dimensions(); + let resized_image_array = resized_image.to_rgb8().into_raw(); + let resized_image_array = Tensor::from_vec( + resized_image_array, + (height as usize, width as usize, 3), + device, + )?; + Ok(resized_image_array) + } + + pub fn split_image( + &self, + image: &DynamicImage, + max_image_size: &HashMap, + resample: FilterType, + device: &Device, + ) -> Result<(Vec, i32, i32), anyhow::Error> { + let (width, height) = image.dimensions(); + let max_size = max_image_size.get("longest_edge").unwrap_or(&364); + let max_height = *max_size; + let max_width = *max_size; + + let mut frames = Vec::new(); + let (num_splits_h, num_splits_w) = if height > max_height as u32 || width > max_width as u32 + { + // Calculate the number of splits + let num_splits_h = (height as f32 / max_height as f32).ceil() as i32; + let num_splits_w = (width as f32 / max_width as f32).ceil() as i32; + + // Calculate optimal dimensions for sub-images + let optimal_height = (height as f32 / num_splits_h as f32).ceil() as u32; + let optimal_width = (width as f32 / num_splits_w as f32).ceil() as u32; + + // Iterate through each row and column + for r in 0..num_splits_h { + for c in 0..num_splits_w { + // Calculate crop coordinates + let start_x = (c as u32) * optimal_width; + let start_y = (r as u32) * optimal_height; + let end_x = std::cmp::min(start_x + optimal_width, width); + let end_y = std::cmp::min(start_y + optimal_height, height); + // Crop the image + let cropped_image = + image.crop_imm(start_x, start_y, end_x - start_x, end_y - start_y); + frames.push(cropped_image); + } + } + + // For the global image at the end, we resize it to match the max_image_size, for cpu memory efficiency + let global_image_height = max_height as u32; + let global_image_width = max_width as u32; + let global_image = if height != global_image_height || width != global_image_width { + let size = HashMap::from([ + ("height".to_string(), global_image_height as i32), + ("width".to_string(), global_image_width as i32), + ]); + let resized = self.resize(image, size, resample, device)?; + DynamicImage::ImageRgb8( + RgbImage::from_raw( + resized.dims()[1] as u32, + resized.dims()[0] as u32, + resized.flatten_all()?.to_vec1::()?, + ) + .ok_or_else(|| { + anyhow::anyhow!("Failed to convert resized image to DynamicImage") + })?, + ) + } else { + image.clone() + }; + frames.push(global_image); + + (num_splits_h, num_splits_w) + } else { + // If image is smaller than max_size, just add it as is + frames.push(image.clone()); + (0, 0) + }; + + let frames = frames + .iter() + .map(|frame| -> Result { + let image = frame.to_rgb8().into_raw(); + Ok(Tensor::from_vec( + image, + (frame.height() as usize, frame.width() as usize, 3), + device, + )?) + }) + .collect::, _>>()?; + + Ok((frames, num_splits_h, num_splits_w)) + } + + fn rescale(&self, image: &Tensor, rescale_factor: f32) -> Result { + let image = (image.to_dtype(DType::F32)? * rescale_factor as f64)?; + Ok(image) + } + + fn pad_image( + &self, + image: &Tensor, + output_size: (i32, i32), + data_format: &str, + device: &Device, + ) -> Result<(Tensor, Tensor), anyhow::Error> { + let (input_height, input_width, _) = image.dims3()?; + let (output_height, output_width) = output_size; + + // Create padded image with zeros + let mut padded_image = if data_format == "channels_first" { + Tensor::zeros( + (3, output_height as usize, output_width as usize), + DType::F32, + device, + )? + } else { + Tensor::zeros( + (output_height as usize, output_width as usize, 3), + DType::F32, + device, + )? + }; + + // Create pixel attention mask + let mut pixel_mask = Tensor::zeros( + (output_height as usize, output_width as usize), + DType::I64, + device, + )?; + + // Copy the original image into the padded image + if data_format == "channels_first" { + let transposed_image = image.to_owned().permute([2, 0, 1])?; + padded_image = padded_image + .slice_assign(&[0..3, 0..input_height, 0..input_width], &transposed_image)?; + pixel_mask = pixel_mask.slice_assign( + &[0..input_height, 0..input_width], + &Tensor::ones((input_height, input_width), DType::I64, device)?, + )?; + } else { + padded_image = + padded_image.slice_assign(&[0..3, 0..input_height, 0..input_width], image)?; + pixel_mask = pixel_mask.slice_assign( + &[0..input_height, 0..input_width], + &Tensor::ones((input_height, input_width), DType::I64, device)?, + )?; + } + + Ok((padded_image, pixel_mask)) + } + + pub fn pad( + &self, + images: &[Vec], + return_pixel_mask: bool, + data_format: &str, + device: &Device, + ) -> Result<(Vec, Option>), anyhow::Error> { + // Get max dimensions across all images + let mut max_height = 0; + let mut max_width = 0; + let mut max_num_images = 0; + + for batch in images { + max_num_images = std::cmp::max(max_num_images, batch.len()); + for image in batch { + let (height, width, _) = image + .dims3() + .map_err(|e| anyhow::anyhow!("Error getting dimensions: {}", e))?; + max_height = std::cmp::max(max_height, height); + max_width = std::cmp::max(max_width, width); + } + } + + let output_size = (max_height as i32, max_width as i32); + let batch_size = images.len(); + + // Create empty padded images and masks + let mut padded_images = vec![ + vec![ + if data_format == "channels_first" { + Tensor::zeros((3, max_height, max_width), DType::F32, device)? + } else { + Tensor::zeros((max_height, max_width, 3), DType::F32, device)? + }; + max_num_images + ]; + batch_size + ]; + + let mut padded_masks = if return_pixel_mask { + Some(vec![ + vec![ + Tensor::zeros( + (max_height, max_width), + DType::I64, + device + )?; + max_num_images + ]; + batch_size + ]) + } else { + None + }; + + // Pad each image + for (batch_idx, batch) in images.iter().enumerate() { + for (sample_idx, image) in batch.iter().enumerate() { + let (padded_image, pixel_mask) = + self.pad_image(image, output_size, data_format, device)?; + padded_images[batch_idx][sample_idx] = padded_image; + if let Some(ref mut masks) = padded_masks { + masks[batch_idx][sample_idx] = pixel_mask; + } + } + } + + Ok(( + padded_images.into_iter().flatten().collect(), + padded_masks.map(|masks| masks.into_iter().flatten().collect()), + )) + } + + fn _preprocess_one_image(&self, image: &DynamicImage, device: &Device) -> ProcessedImageResult { + // Step 1: Initial resize + let resized_image = self.resize( + image, + self.size + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing size"))?, + FilterType::Lanczos3, + device, + )?; + + // Convert back to DynamicImage for further processing + let resized_dynamic_image = DynamicImage::ImageRgb8( + RgbImage::from_raw( + resized_image.dims()[1] as u32, + resized_image.dims()[0] as u32, + resized_image.flatten_all()?.to_vec1::()?, + ) + .ok_or_else(|| anyhow::anyhow!("Failed to convert resized image to DynamicImage"))?, + ); + + // Step 2: Resize for vision encoder + let vision_encoder_image = self.resize_for_vision_encoder( + &resized_dynamic_image, + &self + .max_image_size + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing max_image_size"))?["longest_edge"], + FilterType::Lanczos3, + device, + )?; + + let vision_encoder_image = DynamicImage::ImageRgb8( + RgbImage::from_raw( + vision_encoder_image.dims()[1] as u32, + vision_encoder_image.dims()[0] as u32, + vision_encoder_image.flatten_all()?.to_vec1::()?, + ) + .ok_or_else(|| anyhow::anyhow!("Failed to create RgbImage from raw data"))?, + ); + + // Step 3: Split image if needed + let (frames, n_rows, n_cols) = if self.do_image_splitting { + self.split_image( + &vision_encoder_image, + &self + .max_image_size + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing max_image_size"))?, + FilterType::Lanczos3, + device, + )? + } else { + let frame = vision_encoder_image.to_rgb8().into_raw(); + let frame = Tensor::from_vec( + frame, + ( + vision_encoder_image.height() as usize, + vision_encoder_image.width() as usize, + 3, + ), + device, + )?; + (vec![frame], 1, 1) + }; + + // Step 4 & 5: Rescale and normalize frames + let normalized_frames = if self.do_normalize { + let image_mean = self + .image_mean + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing image_mean"))?; + let image_std = self + .image_std + .clone() + .ok_or_else(|| anyhow::anyhow!("Missing image_std"))?; + let image_mean = + Tensor::from_vec(image_mean.clone(), (1, 1, image_mean.len()), device)?; + let image_std = Tensor::from_vec(image_std.clone(), (1, 1, image_std.len()), device)?; + + frames + .iter() + .map(|frame| { + // First rescale the frame + let rescaled_frame = if self.do_rescale { + self.rescale(frame, self.rescale_factor)? + } else { + frame.to_dtype(DType::F32)? + }; + // Then normalize the rescaled frame + normalize(&rescaled_frame, &image_mean, &image_std) + }) + .collect::, _>>()? + } else { + // Only rescale without normalization + frames + .iter() + .map(|frame| -> Result { + if self.do_rescale { + self.rescale(frame, self.rescale_factor) + } else { + Ok(frame.to_dtype(DType::F32)?) + } + }) + .collect::, _>>()? + }; + + // Step 6: Pad and stack frames + let (padded_images, padded_masks) = if self.do_pad { + self.pad(&[normalized_frames], true, "channels_first", device)? + } else { + (normalized_frames, None) + }; + + Ok((padded_images, padded_masks, n_rows, n_cols)) + } + + pub fn preprocess( + &self, + images: &[DynamicImage], + device: &Device, + ) -> Result<(Tensor, Option, i32, i32), anyhow::Error> { + let mut preprocessed_images = Vec::new(); + let mut preprocessed_masks = Vec::new(); + let mut n_rows = 0; + let mut n_cols = 0; + for image in images { + let (padded_images, padded_masks, rows, cols) = + self._preprocess_one_image(image, device)?; + preprocessed_images.push(padded_images); + if let Some(mask) = padded_masks { + preprocessed_masks.push(mask); + } + n_rows = rows; + n_cols = cols; + } + let max_num_images = preprocessed_images + .iter() + .map(|image| image.len()) + .max() + .expect("No images found"); + let (max_height, max_width) = get_max_height_width(&preprocessed_images); + let mut padded_image_list_full = vec![]; + let mut padded_mask_list_full = vec![]; + for _ in images { + let (padded_image_list, padded_mask_list): (Vec<_>, Vec<_>) = (0..max_num_images) + .map(|_| -> Result<(Tensor, Tensor), anyhow::Error> { + Ok(( + empty_image(3, max_height as u32, max_width as u32, DType::F32, device)?, + Tensor::zeros((max_height, max_width), DType::I64, device)?, + )) + }) + .collect::, _>>()? + .into_iter() + .unzip(); + padded_image_list_full.push(padded_image_list); + padded_mask_list_full.push(padded_mask_list); + } + for (i, image) in preprocessed_images.iter().enumerate() { + padded_image_list_full[i][..image.len()].clone_from_slice(&image[..]); + } + + let (padded_image_list_full, padded_mask_list_full): (Vec<_>, Vec<_>) = + padded_image_list_full + .iter() + .zip(padded_mask_list_full.iter()) + .map( + |(image_list, mask_list)| -> Result<(Tensor, Tensor), anyhow::Error> { + Ok((Tensor::stack(image_list, 0)?, Tensor::stack(mask_list, 0)?)) + }, + ) + .collect::, _>>()? + .into_iter() + .unzip(); + + let padded_image_list_full = Tensor::stack(&padded_image_list_full, 0)?; + let padded_mask_list_full = Tensor::stack(&padded_mask_list_full, 0)?; + + Ok(( + padded_image_list_full, + Some(padded_mask_list_full), + n_rows, + n_cols, + )) + } + + pub fn from_pretrained(model_id: &str) -> Result { + let api = Api::new()?; + let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model)); + let config_file = repo.get("preprocessor_config.json")?; + let processor: Idefics3ImageProcessor = + serde_json::from_slice(&std::fs::read(config_file)?) + .map_err(|e| anyhow::anyhow!("Failed to read preprocessor config: {}", e))?; + Ok(processor) + } +} + +pub struct Idefics3Processor { + image_processor: Idefics3ImageProcessor, + tokenizer: Tokenizer, + fake_image_token: AddedToken, + image_token: AddedToken, + global_image_tag: AddedToken, + image_seq_len: i32, +} + +impl Idefics3Processor { + pub fn from_pretrained(model_id: &str) -> anyhow::Result { + let image_processor = Idefics3ImageProcessor::from_pretrained(model_id)?; + let api = Api::new()?; + let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model)); + + let processor_config_file = repo.get("processor_config.json")?; + let processor_config: serde_json::Value = + serde_json::from_slice(&std::fs::read(processor_config_file)?)?; + let image_seq_len = processor_config["image_seq_len"].as_u64().unwrap_or(64) as i32; + let config_file = repo.get("tokenizer.json")?; + let mut tokenizer = Tokenizer::from_file(config_file) + .map_err(|e| anyhow::anyhow!("Tokenizer error: {}", e))?; + let fake_image_token = AddedToken::from("", true); + let image_token = AddedToken::from("", true); + let end_of_utterance_token = AddedToken::from("", true); + let global_image_tag = AddedToken::from("", true); + tokenizer.add_special_tokens(&[ + fake_image_token.clone(), + image_token.clone(), + end_of_utterance_token.clone(), + global_image_tag.clone(), + ]); + + Ok(Idefics3Processor { + image_processor, + tokenizer, + fake_image_token, + image_token, + global_image_tag, + image_seq_len, + }) + } + + pub fn tokenize_batch( + &self, + prompts: Vec<&str>, + device: &Device, + ) -> Result<(Tensor, Tensor), anyhow::Error> { + let new_prompts = prompts + .iter() + .map(|prompt| { + let prompt = format!("Query: {} {} \n", prompt, "".repeat(10)); + prompt + }) + .collect::>(); + + let tokens = self + .tokenizer + .encode_batch(new_prompts, true) + .map_err(|e| anyhow::anyhow!("Tokenizer error: {}", e))?; + let token_ids = tokens + .iter() + .map(|tokens| { + let tokens = tokens.get_ids().to_vec(); + Tensor::new(tokens.as_slice(), device) + }) + .collect::, _>>()?; + let input = Tensor::stack(&token_ids, 0)?; + let attention_mask = tokens + .iter() + .map(|tokens| { + let tokens = tokens.get_attention_mask().to_vec(); + Tensor::new(tokens.as_slice(), device) + }) + .collect::, _>>()?; + let attention_mask = Tensor::stack(&attention_mask, 0)?; + Ok((input, attention_mask)) + } + + pub fn preprocess( + &self, + images: &[DynamicImage], + device: &Device, + ) -> anyhow::Result<(Tensor, Tensor, Tensor, Option)> { + let (preprocessed_images, preprocessed_masks, n_rows, n_cols) = + self.image_processor.preprocess(images, device)?; + let image_prompt = _prompt_split_image( + self.image_seq_len, + n_rows, + n_cols, + &self.fake_image_token, + &self.image_token, + &self.global_image_tag, + ); + + let prompt = "<|im_start|>user\nDescribe the image."; + + // in the prompt replace the image_token with the image_prompt + let mut prompts = Vec::new(); + for _ in images { + let prompt = prompt.replace(&self.image_token.content, &image_prompt); + prompts.push(prompt); + } + + let (input_ids, attention_mask) = + self.tokenize_batch(prompts.iter().map(|x| x.as_str()).collect(), device)?; + + Ok(( + input_ids, + attention_mask, + preprocessed_images, + preprocessed_masks, + )) + } +} + +fn get_max_height_width(image_list: &[Vec]) -> (usize, usize) { + let mut max_height = 0; + let mut max_width = 0; + for images in image_list { + for image in images { + let height = image.dims()[1]; + let width = image.dims()[2]; + max_height = std::cmp::max(max_height, height); + max_width = std::cmp::max(max_width, width); + } + } + (max_height, max_width) +} + +fn empty_image( + channels: usize, + height: u32, + width: u32, + dtype: DType, + device: &Device, +) -> Result { + Ok(Tensor::zeros( + (channels, height as usize, width as usize), + dtype, + device, + )?) +} + +fn _prompt_split_image( + image_seq_len: i32, + image_rows: i32, + image_cols: i32, + fake_token_around_image: &AddedToken, + image_token: &AddedToken, + global_img_token: &AddedToken, +) -> String { + let mut text_split_images = String::new(); + for n_h in 0..image_rows { + for n_w in 0..image_cols { + text_split_images.push_str(&format!( + "{}{}", + fake_token_around_image.content, + n_h + 1, + n_w + 1, + image_token.content.repeat(image_seq_len as usize), + )); + } + text_split_images.push('\n'); + } + text_split_images.push_str(&format!( + "\n{}{}{}{}", + fake_token_around_image.content, + global_img_token.content, + image_token.content.repeat(image_seq_len as usize), + fake_token_around_image.content + )); + text_split_images +} + +fn get_resize_output_image_size(image: DynamicImage, resolution_max_side: i32) -> (i32, i32) { + let (width, height) = image.dimensions(); + let (new_height, new_width) = + _resize_output_size_rescale_to_max_len(height as i32, width as i32, 1, resolution_max_side); + let (new_height, new_width) = + _resize_output_size_scale_below_upper_bound(new_height, new_width, Some(MAX_IMAGE_SIZE)); + (new_height, new_width) +} + +fn _resize_output_size_rescale_to_max_len( + height: i32, + width: i32, + min_len: i32, + max_len: i32, +) -> (i32, i32) { + let max_len = if max_len == 0 { + std::cmp::max(height, width) + } else { + max_len + }; + let aspect_ratio = width as f32 / height as f32; + let (new_width, new_height) = if width >= height { + let new_width = max_len; + let new_height = (new_width as f32 / aspect_ratio) as i32; + if new_height % 2 != 0 { + (new_width, new_height + 1) + } else { + (new_width, new_height) + } + } else { + let new_height = max_len; + let new_width = (new_height as f32 * aspect_ratio) as i32; + if new_width % 2 != 0 { + (new_width + 1, new_height) + } else { + (new_width, new_height) + } + }; + + // Avoid resizing to a size smaller than min_len + let new_height = std::cmp::max(new_height, min_len); + let new_width = std::cmp::max(new_width, min_len); + (new_height, new_width) +} + +fn _resize_output_size_scale_below_upper_bound( + height: i32, + width: i32, + max_len: Option, +) -> (i32, i32) { + let max_len = max_len.unwrap_or_else(|| std::cmp::max(height, width)); + let aspect_ratio = width as f32 / height as f32; + + let (new_width, new_height) = if width >= height && width > max_len { + let new_width = max_len; + let new_height = (new_width as f32 / aspect_ratio) as i32; + (new_width, new_height) + } else if height > width && height > max_len { + let new_height = max_len; + let new_width = (new_height as f32 * aspect_ratio) as i32; + (new_width, new_height) + } else { + (width, height) + }; + + // Avoid resizing to a size smaller than 1 + let new_height = std::cmp::max(new_height, 1); + let new_width = std::cmp::max(new_width, 1); + (new_height, new_width) +} + +fn normalize(image: &Tensor, mean: &Tensor, std: &Tensor) -> Result { + let normalized_image = image + .to_dtype(DType::F32)? + .broadcast_sub(mean)? + .broadcast_div(std)?; + Ok(normalized_image) +} diff --git a/rust/src/models/llama.rs b/rust/src/models/llama.rs new file mode 100644 index 00000000..86a215c9 --- /dev/null +++ b/rust/src/models/llama.rs @@ -0,0 +1,587 @@ +//! Llama inference implementation. +//! +//! See ["LLaMA: Open and Efficient Foundation Language Models"](https://arxiv.org/abs/2302.13971) +//! +//! Implementation based on Hugging Face's [transformers](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py) + +use super::with_tracing::{linear_no_bias as linear, Linear, RmsNorm}; +use candle_core::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{embedding, Embedding, Module, VarBuilder}; +use std::{collections::HashMap, f32::consts::PI}; + +pub const DEFAULT_MAX_SEQ_LEN: usize = 4096; + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub enum Llama3RopeType { + #[serde(rename = "llama3")] + Llama3, + #[default] + #[serde(rename = "default")] + Default, +} + +#[derive(Debug, Clone, serde::Deserialize, Default)] +pub struct Llama3RopeConfig { + pub factor: f32, + pub low_freq_factor: f32, + pub high_freq_factor: f32, + pub original_max_position_embeddings: usize, + pub rope_type: Llama3RopeType, +} +#[derive(Debug, Clone, serde::Deserialize)] +#[serde(untagged)] +pub enum LlamaEosToks { + Single(u32), + Multiple(Vec), +} + +#[derive(Debug, Clone, serde::Deserialize)] +pub struct LlamaConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: Option, + pub rms_norm_eps: f64, + #[serde(default = "default_rope")] + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub tie_word_embeddings: Option, +} + +impl LlamaConfig { + pub fn num_key_value_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } +} + +fn default_rope() -> f32 { + 10_000.0 +} + +impl LlamaConfig { + pub fn into_config(self, use_flash_attn: bool) -> Config { + Config { + hidden_size: self.hidden_size, + intermediate_size: self.intermediate_size, + vocab_size: self.vocab_size, + num_hidden_layers: self.num_hidden_layers, + num_attention_heads: self.num_attention_heads, + num_key_value_heads: self.num_key_value_heads(), + rms_norm_eps: self.rms_norm_eps, + rope_theta: self.rope_theta, + use_flash_attn, + bos_token_id: self.bos_token_id, + eos_token_id: self.eos_token_id, + rope_scaling: self.rope_scaling, + max_position_embeddings: self.max_position_embeddings, + tie_word_embeddings: self.tie_word_embeddings.unwrap_or(false), + } + } +} + +#[derive(Debug, Clone)] +pub struct Config { + pub hidden_size: usize, + pub intermediate_size: usize, + pub vocab_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub use_flash_attn: bool, + pub rms_norm_eps: f64, + pub rope_theta: f32, + pub bos_token_id: Option, + pub eos_token_id: Option, + pub rope_scaling: Option, + pub max_position_embeddings: usize, + pub tie_word_embeddings: bool, +} + +impl Config { + pub fn config_7b_v1(use_flash_attn: bool) -> Self { + Self { + hidden_size: 4096, + intermediate_size: 11008, + vocab_size: 32000, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 32, + use_flash_attn, + rms_norm_eps: 1e-6, + rope_theta: 10_000.0, + bos_token_id: None, + eos_token_id: None, + rope_scaling: None, + max_position_embeddings: DEFAULT_MAX_SEQ_LEN, + tie_word_embeddings: false, + } + } + + pub fn config_7b_v2(use_flash_attn: bool) -> Self { + Self { + hidden_size: 4096, + intermediate_size: 11008, + vocab_size: 32000, + num_hidden_layers: 32, + num_attention_heads: 32, + num_key_value_heads: 32, + use_flash_attn, + rms_norm_eps: 1e-5, + rope_theta: 10_000.0, + bos_token_id: None, + eos_token_id: None, + rope_scaling: None, + max_position_embeddings: DEFAULT_MAX_SEQ_LEN, + tie_word_embeddings: false, + } + } +} + +#[derive(Debug, Clone)] +pub struct Cache { + masks: HashMap, + pub use_kv_cache: bool, + kvs: Vec>, + cos: Tensor, + sin: Tensor, + device: Device, +} + +fn calculate_default_inv_freq(cfg: &Config) -> Vec { + let head_dim = cfg.hidden_size / cfg.num_attention_heads; + (0..head_dim) + .step_by(2) + .map(|i| 1f32 / cfg.rope_theta.powf(i as f32 / head_dim as f32)) + .collect() +} + +impl Cache { + pub fn new(use_kv_cache: bool, dtype: DType, config: &Config, device: &Device) -> Result { + // precompute freqs_cis + let theta = match &config.rope_scaling { + None + | Some(Llama3RopeConfig { + rope_type: Llama3RopeType::Default, + .. + }) => calculate_default_inv_freq(config), + Some(rope_scaling) => { + let low_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.low_freq_factor; + let high_freq_wavelen = rope_scaling.original_max_position_embeddings as f32 + / rope_scaling.high_freq_factor; + + calculate_default_inv_freq(config) + .into_iter() + .map(|freq| { + let wavelen = 2. * PI / freq; + if wavelen < high_freq_wavelen { + freq + } else if wavelen > low_freq_wavelen { + freq / rope_scaling.factor + } else { + let smooth = (rope_scaling.original_max_position_embeddings as f32 + / wavelen + - rope_scaling.low_freq_factor) + / (rope_scaling.high_freq_factor - rope_scaling.low_freq_factor); + (1. - smooth) * freq / rope_scaling.factor + smooth * freq + } + }) + .collect::>() + } + }; + + let theta = Tensor::new(theta, device)?; + + let idx_theta = Tensor::arange(0, config.max_position_embeddings as u32, device)? + .to_dtype(DType::F32)? + .reshape((config.max_position_embeddings, 1))? + .matmul(&theta.reshape((1, theta.elem_count()))?)?; + // This is different from the paper, see: + // https://github.com/huggingface/transformers/blob/6112b1c6442aaf7affd2b0676a1cd4eee30c45cf/src/transformers/models/llama/modeling_llama.py#L112 + let cos = idx_theta.cos()?.to_dtype(dtype)?; + let sin = idx_theta.sin()?.to_dtype(dtype)?; + Ok(Self { + masks: HashMap::new(), + use_kv_cache, + kvs: vec![None; config.num_hidden_layers], + device: device.clone(), + cos, + sin, + }) + } + + fn mask(&mut self, t: usize) -> Result { + if let Some(mask) = self.masks.get(&t) { + Ok(mask.clone()) + } else { + let mask: Vec<_> = (0..t) + .flat_map(|i| (0..t).map(move |j| u8::from(j > i))) + .collect(); + let mask = Tensor::from_slice(&mask, (t, t), &self.device)?; + self.masks.insert(t, mask.clone()); + Ok(mask) + } + } +} + +#[derive(Debug, Clone)] +struct CausalSelfAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_attention_heads: usize, + num_key_value_heads: usize, + head_dim: usize, + use_flash_attn: bool, + span: tracing::Span, + span_rot: tracing::Span, + max_position_embeddings: usize, +} + +#[cfg(feature = "flash-attn")] +fn flash_attn( + q: &Tensor, + k: &Tensor, + v: &Tensor, + softmax_scale: f32, + causal: bool, +) -> Result { + candle_flash_attn::flash_attn(q, k, v, softmax_scale, causal) +} + +#[cfg(not(feature = "flash-attn"))] +fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result { + unimplemented!("compile with '--features flash-attn'") +} + +impl CausalSelfAttention { + fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize, cache: &Cache) -> Result { + let _enter = self.span_rot.enter(); + let (_b_sz, _, seq_len, _hidden_size) = x.dims4()?; + let cos = cache.cos.narrow(0, index_pos, seq_len)?; + let sin = cache.sin.narrow(0, index_pos, seq_len)?; + candle_nn::rotary_emb::rope(x, &cos, &sin) + } + + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let _enter = self.span.enter(); + let (b_sz, seq_len, hidden_size) = x.dims3()?; + let q = self.q_proj.forward(x)?; + let k = self.k_proj.forward(x)?; + let v = self.v_proj.forward(x)?; + + let q = q + .reshape((b_sz, seq_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let k = k + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)? + .contiguous()?; + let mut v = v + .reshape((b_sz, seq_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + + let q = self.apply_rotary_emb(&q, index_pos, cache)?; + let mut k = self.apply_rotary_emb(&k, index_pos, cache)?; + + if cache.use_kv_cache { + if let Some((cache_k, cache_v)) = &cache.kvs[block_idx] { + k = Tensor::cat(&[cache_k, &k], 2)?.contiguous()?; + v = Tensor::cat(&[cache_v, &v], 2)?.contiguous()?; + let k_seq_len = k.dims()[1]; + if k_seq_len > self.max_position_embeddings { + k = k + .narrow( + D::Minus1, + k_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + let v_seq_len = v.dims()[1]; + if v_seq_len > 2 * self.max_position_embeddings { + v = v + .narrow( + D::Minus1, + v_seq_len - self.max_position_embeddings, + self.max_position_embeddings, + )? + .contiguous()? + } + } + cache.kvs[block_idx] = Some((k.clone(), v.clone())) + } + + let k = self.repeat_kv(k)?; + let v = self.repeat_kv(v)?; + + let y = if self.use_flash_attn { + // flash-attn expects (b_sz, seq_len, nheads, head_dim) + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let softmax_scale = 1f32 / (self.head_dim as f32).sqrt(); + flash_attn(&q, &k, &v, softmax_scale, seq_len > 1)?.transpose(1, 2)? + } else { + let in_dtype = q.dtype(); + let q = q.to_dtype(DType::F32)?; + let k = k.to_dtype(DType::F32)?; + let v = v.to_dtype(DType::F32)?; + let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?; + let att = if seq_len == 1 { + att + } else { + let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; + masked_fill(&att, &mask, f32::NEG_INFINITY)? + }; + + let att = candle_nn::ops::softmax_last_dim(&att)?; + // Convert to contiguous as matmul doesn't support strided vs for now. + att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)? + }; + let y = y.transpose(1, 2)?.reshape(&[b_sz, seq_len, hidden_size])?; + let y = self.o_proj.forward(&y)?; + Ok(y) + } + + fn repeat_kv(&self, x: Tensor) -> Result { + candle_transformers::utils::repeat_kv( + x, + self.num_attention_heads / self.num_key_value_heads, + ) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "attn"); + let span_rot = tracing::span!(tracing::Level::TRACE, "attn-rot"); + let size_in = cfg.hidden_size; + let size_q = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_attention_heads; + let size_kv = (cfg.hidden_size / cfg.num_attention_heads) * cfg.num_key_value_heads; + let q_proj = linear(size_in, size_q, vb.pp("q_proj"))?; + let k_proj = linear(size_in, size_kv, vb.pp("k_proj"))?; + let v_proj = linear(size_in, size_kv, vb.pp("v_proj"))?; + let o_proj = linear(size_q, size_in, vb.pp("o_proj"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_attention_heads: cfg.num_attention_heads, + num_key_value_heads: cfg.num_key_value_heads, + head_dim: cfg.hidden_size / cfg.num_attention_heads, + use_flash_attn: cfg.use_flash_attn, + span, + span_rot, + max_position_embeddings: cfg.max_position_embeddings, + }) + } +} + +fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result { + let shape = mask.shape(); + let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?; + let m = mask.where_cond(&on_true, on_false)?; + Ok(m) +} + +#[derive(Debug, Clone)] +struct Mlp { + c_fc1: Linear, + c_fc2: Linear, + c_proj: Linear, + span: tracing::Span, +} + +impl Mlp { + fn forward(&self, x: &Tensor) -> Result { + let _enter = self.span.enter(); + let x = (candle_nn::ops::silu(&self.c_fc1.forward(x)?)? * self.c_fc2.forward(x)?)?; + self.c_proj.forward(&x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "mlp"); + let h_size = cfg.hidden_size; + let i_size = cfg.intermediate_size; + let c_fc1 = linear(h_size, i_size, vb.pp("gate_proj"))?; + let c_fc2 = linear(h_size, i_size, vb.pp("up_proj"))?; + let c_proj = linear(i_size, h_size, vb.pp("down_proj"))?; + Ok(Self { + c_fc1, + c_fc2, + c_proj, + span, + }) + } +} + +#[derive(Debug, Clone)] +struct Block { + rms_1: RmsNorm, + attn: CausalSelfAttention, + rms_2: RmsNorm, + mlp: Mlp, + span: tracing::Span, +} + +impl Block { + fn forward( + &self, + x: &Tensor, + index_pos: usize, + block_idx: usize, + cache: &mut Cache, + ) -> Result { + let _enter = self.span.enter(); + let residual = x; + let x = self.rms_1.forward(x)?; + let x = (self.attn.forward(&x, index_pos, block_idx, cache)? + residual)?; + let residual = &x; + let x = (self.mlp.forward(&self.rms_2.forward(&x)?)? + residual)?; + Ok(x) + } + + fn load(vb: VarBuilder, cfg: &Config) -> Result { + let span = tracing::span!(tracing::Level::TRACE, "block"); + let attn = CausalSelfAttention::load(vb.pp("self_attn"), cfg)?; + let mlp = Mlp::load(vb.pp("mlp"), cfg)?; + let rms_1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; + let rms_2 = RmsNorm::new( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + )?; + Ok(Self { + rms_1, + attn, + rms_2, + mlp, + span, + }) + } +} + +pub struct LlamaBase { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, +} + +impl LlamaBase { + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?; + let layers: Vec<_> = (0..cfg.num_hidden_layers) + .map(|i| Block::load(vb.pp(format!("layers.{i}")), cfg).unwrap()) + .collect(); + let norm = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?; + Ok(Self { + embed_tokens, + layers, + norm, + }) + } + + pub fn embed(&self, x: &Tensor) -> Result { + self.embed_tokens.forward(x) + } + + pub fn forward_input_embed( + &self, + x: &Tensor, + index_pos: usize, + cache: &mut Cache, + ) -> Result { + let mut x = x.clone(); + for (block_idx, block) in self.layers.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.norm.forward(&x)?; + Ok(x) + } + + pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result { + let mut x = self.embed_tokens.forward(x)?; + for (block_idx, block) in self.layers.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.norm.forward(&x)?; + Ok(x) + } +} + +#[derive(Debug, Clone)] +pub struct Llama { + wte: Embedding, + blocks: Vec, + ln_f: RmsNorm, + lm_head: Linear, +} + +impl Llama { + // required by LLaVA + pub fn embed(&self, x: &Tensor) -> Result { + self.wte.forward(x) + } + // required by LLaVA + pub fn forward_input_embed( + &self, + input_embed: &Tensor, + index_pos: usize, + cache: &mut Cache, + ) -> Result { + let (_, seq_len, _) = input_embed.dims3()?; + let mut x = input_embed.clone(); + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn forward(&self, x: &Tensor, index_pos: usize, cache: &mut Cache) -> Result { + let (_b_sz, seq_len) = x.dims2()?; + let mut x = self.wte.forward(x)?; + for (block_idx, block) in self.blocks.iter().enumerate() { + x = block.forward(&x, index_pos, block_idx, cache)?; + } + let x = self.ln_f.forward(&x)?; + let x = x.i((.., seq_len - 1, ..))?.contiguous()?; + let logits = self.lm_head.forward(&x)?; + logits.to_dtype(DType::F32) + } + + pub fn load(vb: VarBuilder, cfg: &Config) -> Result { + let wte = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?; + let lm_head = if cfg.tie_word_embeddings { + Linear::from_weights(wte.embeddings().clone(), None) + } else { + linear(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))? + }; + let ln_f = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?; + let blocks: Vec<_> = (0..cfg.num_hidden_layers) + .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap()) + .collect(); + + Ok(Self { + wte, + blocks, + ln_f, + lm_head, + }) + } +} diff --git a/rust/src/models/mod.rs b/rust/src/models/mod.rs index 22bbb218..cb72c167 100644 --- a/rust/src/models/mod.rs +++ b/rust/src/models/mod.rs @@ -2,10 +2,12 @@ pub mod bert; pub mod clip; pub mod colpali; pub mod gemma; +pub mod idefics3; pub mod jina_bert; +pub mod llama; pub mod modernbert; pub mod paligemma; +pub mod quantized_qwen3; +pub mod qwen3; pub mod siglip; pub mod with_tracing; -pub mod qwen3; -pub mod quantized_qwen3; \ No newline at end of file diff --git a/rust/src/models/quantized_qwen3.rs b/rust/src/models/quantized_qwen3.rs index 039e838f..5491e114 100644 --- a/rust/src/models/quantized_qwen3.rs +++ b/rust/src/models/quantized_qwen3.rs @@ -7,14 +7,13 @@ //! - [Qwen3 Models](https://huggingface.co/Qwen/Qwen3-0.6B) (architecture based on official implementations) //! use super::with_tracing::QMatMul; -use candle_transformers::{quantized_nn::RmsNorm, utils::repeat_kv}; use candle_core::quantized::{gguf_file, QTensor}; use candle_core::{DType, Device, Result, Tensor}; use candle_nn::{kv_cache::KvCache, Activation, Embedding, Module}; +use candle_transformers::{quantized_nn::RmsNorm, utils::repeat_kv}; use std::io::{Read, Seek}; use std::sync::Arc; - struct Gguf { ct: gguf_file::Content, reader: R, @@ -427,4 +426,4 @@ impl ModelWeights { let last_hidden = h.narrow(1, l - 1, 1)?; self.lm_head.forward(&last_hidden)?.squeeze(1) } -} \ No newline at end of file +} diff --git a/rust/src/models/qwen3.rs b/rust/src/models/qwen3.rs index bf6dd1e4..fe2922b0 100644 --- a/rust/src/models/qwen3.rs +++ b/rust/src/models/qwen3.rs @@ -401,5 +401,5 @@ fn prepare_4d_attention_mask( let inverted_mask = (1.0 - expanded_mask)?; - (inverted_mask * -1e4 as f64)?.to_dtype(dtype) + (inverted_mask * -1e4_f64)?.to_dtype(dtype) } diff --git a/rust/src/models/siglip.rs b/rust/src/models/siglip.rs index 9a1fe642..65ddd3d6 100644 --- a/rust/src/models/siglip.rs +++ b/rust/src/models/siglip.rs @@ -795,4 +795,4 @@ impl Model { let logits_per_image = logits_per_text.t()?; Ok((logits_per_text, logits_per_image)) } -} \ No newline at end of file +} diff --git a/rust/src/reranker/model.rs b/rust/src/reranker/model.rs index 768f206e..cd857b41 100644 --- a/rust/src/reranker/model.rs +++ b/rust/src/reranker/model.rs @@ -118,7 +118,10 @@ impl Reranker { .with_inter_threads(1)? // Set inter-op parallelism to 1 when using GPU .commit_from_file(weights_filename)?; - Ok(Reranker { model: RwLock::new(model), tokenizer }) + Ok(Reranker { + model: RwLock::new(model), + tokenizer, + }) } pub fn compute_scores( diff --git a/rust/src/text_loader.rs b/rust/src/text_loader.rs index e4b51c65..dc2ad047 100644 --- a/rust/src/text_loader.rs +++ b/rust/src/text_loader.rs @@ -1,8 +1,4 @@ -use std::{ - collections::HashMap, - fmt::Debug, - fs, -}; +use std::{collections::HashMap, fmt::Debug, fs}; use anyhow::Error; use chrono::{DateTime, Local};