diff --git a/rust/src/chunkers/mod.rs b/rust/src/chunkers/mod.rs index 09bbef29..4d1624e1 100644 --- a/rust/src/chunkers/mod.rs +++ b/rust/src/chunkers/mod.rs @@ -1,2 +1,20 @@ +//! Text chunking algorithm implementations. +//! +//! Provides low-level chunking algorithms that support the text splitting +//! strategies available in the embedding pipeline. These algorithms balance +//! context preservation with computational efficiency. +//! +//! # Available Algorithms +//! +//! - **Statistical** - Length-based chunking with statistical overlap +//! - **Cumulative** - Boundary-aware accumulation chunking +//! +//! # Usage +//! +//! These algorithms are used internally by the embedding pipeline. +//! End users configure chunking through [`TextEmbedConfig`]. +//! +//! [`TextEmbedConfig`]: crate::config::TextEmbedConfig + pub mod cumulative; pub mod statistical; diff --git a/rust/src/config.rs b/rust/src/config.rs index 9bea9a1c..4655519c 100644 --- a/rust/src/config.rs +++ b/rust/src/config.rs @@ -1,3 +1,8 @@ +//! Configuration structs and enums for embedding operations. +//! +//! Provides configuration options for text and image embedding processes, +//! including chunking strategies, batch sizes, and splitting methods. + use processors_rs::pdf::pdf_processor::PdfBackend; use crate::embeddings::embed::Embedder; diff --git a/rust/src/embeddings/cloud/mod.rs b/rust/src/embeddings/cloud/mod.rs index 09597487..dacaea74 100644 --- a/rust/src/embeddings/cloud/mod.rs +++ b/rust/src/embeddings/cloud/mod.rs @@ -1,3 +1,8 @@ +//! Cloud-based embedding service integrations. +//! +//! Implementations for embedding models provided by cloud services. +//! Requires API keys and internet connectivity. + pub mod cohere; pub mod gemini; pub mod openai; diff --git a/rust/src/embeddings/local/mod.rs b/rust/src/embeddings/local/mod.rs index 43219a27..d071412b 100644 --- a/rust/src/embeddings/local/mod.rs +++ b/rust/src/embeddings/local/mod.rs @@ -1,3 +1,8 @@ +//! Local embedding model implementations. +//! +//! Self-contained embedding models that run locally without external API calls. +//! Models use either Candle backend or ONNX Runtime for inference. + pub mod bert; pub mod clip; #[cfg(feature = "ort")] diff --git a/rust/src/embeddings/mod.rs b/rust/src/embeddings/mod.rs index 1ed4ed5d..1b252016 100644 --- a/rust/src/embeddings/mod.rs +++ b/rust/src/embeddings/mod.rs @@ -1,4 +1,7 @@ -//! This module contains the different embedding models that can be used to generate embeddings for the text data. +//! Embedding model implementations and utilities. +//! +//! Provides local and cloud-based embedding models for generating vector +//! representations from text, image, and audio data. use std::{collections::HashMap, rc::Rc}; diff --git a/rust/src/file_loader.rs b/rust/src/file_loader.rs index dc0c4d45..132c5fa4 100644 --- a/rust/src/file_loader.rs +++ b/rust/src/file_loader.rs @@ -1,3 +1,8 @@ +//! File discovery and parsing utilities. +//! +//! Provides functionality for discovering and filtering files in directories +//! based on patterns and file types for embedding processing. + use std::{collections::HashSet, io::Error, path::PathBuf}; use regex::Regex; diff --git a/rust/src/file_processor/audio/mod.rs b/rust/src/file_processor/audio/mod.rs index 54de5bd5..f46331e8 100644 --- a/rust/src/file_processor/audio/mod.rs +++ b/rust/src/file_processor/audio/mod.rs @@ -1,2 +1,12 @@ +//! Audio processing implementation modules. +//! +//! Core audio processing functionality for speech-to-text conversion +//! and audio format handling. +//! +//! # Modules +//! +//! - [`audio_processor`] - Main audio processing and transcription pipeline +//! - [`pcm_decode`] - PCM audio format decoding utilities + pub mod audio_processor; pub mod pcm_decode; diff --git a/rust/src/file_processor/mod.rs b/rust/src/file_processor/mod.rs index fa3e4e46..bdfea052 100644 --- a/rust/src/file_processor/mod.rs +++ b/rust/src/file_processor/mod.rs @@ -1 +1,24 @@ +//! Audio file processing utilities. +//! +//! Provides specialized processors for extracting and transcribing audio content +//! with temporal segmentation for embedding generation. +//! +//! # Features +//! +//! - **Transcription** - Speech-to-text conversion +//! - **Temporal segmentation** - Time-based audio chunking +//! - **Multiple formats** - Support for various audio file types +//! +//! # Usage +//! +//! Audio processing is handled through the main embedding functions: +//! +//! ```rust +//! use embed_anything::embed_file; +//! # async fn example() -> anyhow::Result<()> { +//! let embeddings = embed_file("audio.wav", &embedder, None, None).await?; +//! # Ok(()) +//! # } +//! ``` + pub mod audio; diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 62a0bd5d..79132d2a 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -51,14 +51,21 @@ //! Let's see how embed_anything can help us generate embeddings from a plain text file: //! //! ```rust -//! use embed_anything::embed_file; -//! use embed_anything::embeddings::embed::{Embedder, TextEmbedder}; -//! use embed_anything::embeddings::local::jina::JinaEmbedder; +//! use embed_anything::{embed_file, embeddings::embed::EmbedderBuilder}; +//! use embed_anything::config::TextEmbedConfig; //! -//! // Create an Embedder for text. We support a variety of models out-of-the-box, including cloud-based models! -//! let embedder = Embedder::Text(TextEmbedder::Jina(Box::new(JinaEmbedder::default()))); -//! // Generate embeddings for 'path/to/file.txt' using the embedder we just created. -//! let embedding = embed_file("path/to/file.txt", &embedder, None, None); +//! # async fn example() -> anyhow::Result<()> { +//! // Create an embedder using a pre-trained model from Hugging Face +//! let embedder = EmbedderBuilder::new() +//! .model_architecture("jina") +//! .model_id(Some("jinaai/jina-embeddings-v2-small-en")) +//! .from_pretrained_hf()?; +//! let config = TextEmbedConfig::default(); +//! +//! // Generate embeddings for any supported file type +//! let embeddings = embed_file("document.pdf", &embedder, Some(&config), None).await?; +//! # Ok(()) +//! # } //! ``` pub mod chunkers; @@ -97,15 +104,25 @@ use processors_rs::{ txt_processor::TxtProcessor, }; +/// Numerical precision types for model weights and computations. pub enum Dtype { + /// 16-bit floating point. F16, + /// 8-bit signed integer. INT8, + /// 4-bit quantized. Q4, + /// 8-bit unsigned integer. UINT8, + /// 4-bit BitsAndBytes quantization. BNB4, + /// 32-bit floating point. F32, + /// 4-bit quantized with 16-bit float scale. Q4F16, + /// Generic quantized format. QUANTIZED, + /// 16-bit brain floating point. BF16, } @@ -128,15 +145,18 @@ pub enum Dtype { /// /// # Example /// -/// ``` -/// use embed_anything::embed_query; -/// use embed_anything::embeddings::embed::{Embedder, TextEmbedder}; -/// use embed_anything::embeddings::local::jina::JinaEmbedder; -/// -/// let query = vec!["Hello".to_string(), "World".to_string()]; -/// let embedder = Embedder::Text(TextEmbedder::Jina(Box::new(JinaEmbedder::default()))); -/// let embeddings = embed_query(query, &embedder, None).await().unwrap(); -/// println!("{:?}", embeddings); +/// ```rust +/// use embed_anything::{embed_query, embeddings::embed::EmbedderBuilder}; +/// +/// # async fn example() -> anyhow::Result<()> { +/// let embedder = EmbedderBuilder::new() +/// .model_architecture("jina") +/// .model_id(Some("jinaai/jina-embeddings-v2-small-en")) +/// .from_pretrained_hf()?; +/// let queries = ["Hello world", "How are you?"]; +/// let embeddings = embed_query(&queries, &embedder, None).await?; +/// # Ok(()) +/// # } /// ``` pub async fn embed_query( query: &[&str], @@ -175,13 +195,18 @@ pub async fn embed_query( /// # Example /// /// ```rust -/// use embed_anything::embed_file; -/// use embed_anything::embeddings::embed::{Embedder, TextEmbedder}; -/// use embed_anything::embeddings::local::bert::BertEmbedder; +/// use embed_anything::{embed_file, embeddings::embed::EmbedderBuilder}; +/// use embed_anything::config::TextEmbedConfig; /// -/// let file_name = "path/to/file.pdf"; -/// let embedder = Embedder::Text(TextEmbedder::from(BertEmbedder::new("sentence-transformers/all-MiniLM-L12-v2".into(), None, None).unwrap())); -/// let embeddings = embed_file(file_name, &embedder, None, None).unwrap(); +/// # async fn example() -> anyhow::Result<()> { +/// let embedder = EmbedderBuilder::new() +/// .model_architecture("bert") +/// .model_id(Some("sentence-transformers/all-MiniLM-L12-v2")) +/// .from_pretrained_hf()?; +/// let config = TextEmbedConfig::default(); +/// let embeddings = embed_file("document.pdf", &embedder, Some(&config), None).await?; +/// # Ok(()) +/// # } /// ``` pub async fn embed_file>( file_name: T, @@ -224,13 +249,22 @@ pub async fn embed_file>( /// /// # Example /// -/// ``` -/// use embed_anything::embed_webpage; -/// use embed_anything::embeddings::embed::{Embedder, TextEmbedder}; -/// use embed_anything::embeddings::local::jina::JinaEmbedder; +/// ```rust +/// use embed_anything::{embed_webpage, embeddings::embed::EmbedderBuilder}; +/// use embed_anything::config::TextEmbedConfig; /// -/// let embedder = Embedder::Text(TextEmbedder::Jina(Box::new(JinaEmbedder::default()))); -/// let embeddings = embed_webpage("https://en.wikipedia.org/wiki/Embedding".into(), &embedder, None, None).unwrap(); +/// # async fn example() -> anyhow::Result<()> { +/// let embedder = EmbedderBuilder::new() +/// .model_architecture("jina") +/// .model_id(Some("jinaai/jina-embeddings-v2-small-en")) +/// .from_pretrained_hf()?; +/// let config = TextEmbedConfig::default(); +/// let embeddings = embed_webpage( +/// "https://example.com".to_string(), +/// &embedder, Some(&config), None +/// ).await?; +/// # Ok(()) +/// # } /// ``` pub async fn embed_webpage( url: String, @@ -331,7 +365,24 @@ async fn emb_image>( Ok(embedding) } +/// Embeds audio content from a file using transcription and temporal segmentation. +/// +/// Processes audio files by first transcribing them to text and then creating +/// embeddings from the transcribed segments with temporal information. +/// +/// # Arguments +/// +/// * `audio_file` - Path to the audio file to process +/// * `audio_decoder` - Audio decoder model for transcription +/// * `embedder` - Embedding model for generating vector representations +/// * `text_embed_config` - Optional configuration for text embedding parameters +/// +/// # Returns +/// +/// A vector of `EmbedData` objects containing embeddings for each audio segment, +/// or `None` if no audio content was processed. #[cfg(feature = "audio")] +#[cfg_attr(docsrs, doc(cfg(feature = "audio")))] pub async fn emb_audio>( audio_file: T, audio_decoder: &mut AudioDecoderModel, @@ -904,6 +955,22 @@ pub async fn embed_files_batch( } } +/// Processes text chunks into embeddings with metadata preservation. +/// +/// Takes a collection of text chunks and their associated metadata to generate +/// embeddings using the specified embedding model with configurable batch processing. +/// +/// # Arguments +/// +/// * `chunks` - Array of text strings to embed +/// * `metadata` - Array of optional metadata maps for each chunk +/// * `embedding_model` - Embedding model for generating vector representations +/// * `batch_size` - Optional batch size for processing chunks in groups +/// * `late_chunking` - Optional flag to enable late chunking optimization +/// +/// # Returns +/// +/// An `Arc>` containing embeddings and metadata for all processed chunks. pub async fn process_chunks( chunks: &[String], metadata: &[Option>], @@ -965,9 +1032,44 @@ fn extract_document( } } +/// Errors that can occur during file loading and processing. +/// +/// Provides specific error types for common file processing failures, +/// enabling appropriate error handling and user feedback. +/// +/// # Error Handling +/// +/// ```rust +/// use embed_anything::{embed_file, FileLoadingError}; +/// +/// # async fn example() -> Result<(), Box> { +/// match embed_file("document.xyz", &embedder, None, None).await { +/// Err(e) if e.downcast_ref::().is_some() => { +/// match e.downcast_ref::().unwrap() { +/// FileLoadingError::FileNotFound(path) => { +/// eprintln!("File not found: {}", path); +/// } +/// FileLoadingError::UnsupportedFileType(ext) => { +/// eprintln!("Unsupported format: {}", ext); +/// } +/// } +/// } +/// Ok(embeddings) => { /* process embeddings */ } +/// Err(e) => { /* handle other errors */ } +/// } +/// # Ok(()) +/// # } +/// ``` #[derive(Debug)] pub enum FileLoadingError { + /// File does not exist at the specified path. + /// + /// Common causes: + /// - Incorrect file path + /// - File permissions + /// - File moved or deleted FileNotFound(String), + /// File format is not supported for processing. UnsupportedFileType(String), } impl Display for FileLoadingError { diff --git a/rust/src/models/idefics3/mod.rs b/rust/src/models/idefics3/mod.rs index 0f38153f..7744ef51 100644 --- a/rust/src/models/idefics3/mod.rs +++ b/rust/src/models/idefics3/mod.rs @@ -1,2 +1,6 @@ +//! Idefics3 multimodal model implementation. +//! +//! Vision-language model components for multimodal embedding generation. + pub mod model; pub mod tensor_processing; diff --git a/rust/src/models/mod.rs b/rust/src/models/mod.rs index cb72c167..5ba30277 100644 --- a/rust/src/models/mod.rs +++ b/rust/src/models/mod.rs @@ -1,3 +1,21 @@ +//! Neural network model implementations for embedding generation. +//! +//! Candle-based implementations of embedding models for local inference +//! with optional GPU acceleration support. +//! +//! # Usage +//! +//! Models are accessed through the embedding API: +//! +//! ```rust +//! use embed_anything::embeddings::embed::EmbedderBuilder; +//! +//! let embedder = EmbedderBuilder::new() +//! .model_architecture("bert") +//! .model_id(Some("sentence-transformers/all-MiniLM-L6-v2")) +//! .from_pretrained_hf()?; +//! ``` + pub mod bert; pub mod clip; pub mod colpali; diff --git a/rust/src/reranker/mod.rs b/rust/src/reranker/mod.rs index 18a2a381..901c9263 100644 --- a/rust/src/reranker/mod.rs +++ b/rust/src/reranker/mod.rs @@ -1,2 +1,6 @@ +//! Document reranking model implementations. +//! +//! Models for reordering search results based on relevance scores. + pub mod model; pub mod qwen3; \ No newline at end of file diff --git a/rust/src/text_loader.rs b/rust/src/text_loader.rs index dc2ad047..54404034 100644 --- a/rust/src/text_loader.rs +++ b/rust/src/text_loader.rs @@ -1,3 +1,8 @@ +//! Text loading and chunking utilities. +//! +//! Provides functionality for loading text content from files and splitting +//! it into manageable chunks for embedding generation. + use std::{collections::HashMap, fmt::Debug, fs}; use anyhow::Error;