diff --git a/README.md b/README.md index 9348169..769bc4a 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,7 @@ Quantized versions are also available for several models above (append `Q` to th - [**BAAI/bge-reranker-base**](https://huggingface.co/BAAI/bge-reranker-base) - Default - [**BAAI/bge-reranker-v2-m3**](https://huggingface.co/BAAI/bge-reranker-v2-m3) - [**jinaai/jina-reranker-v1-turbo-en**](https://huggingface.co/jinaai/jina-reranker-v1-turbo-en) -- [**jinaai/jina-reranker-v2-base-multiligual**](https://huggingface.co/jinaai/jina-reranker-v2-base-multilingual) +- [**jinaai/jina-reranker-v2-base-multilingual**](https://huggingface.co/jinaai/jina-reranker-v2-base-multilingual) diff --git a/src/bgem3_embedding/init.rs b/src/bgem3_embedding/init.rs index 89d7afe..503015c 100644 --- a/src/bgem3_embedding/init.rs +++ b/src/bgem3_embedding/init.rs @@ -26,7 +26,7 @@ pub struct Bgem3EmbeddingOutput { /// Struct for "bring your own" BGE-M3 models /// -/// The onnx_file and tokenizer_files expect the files' bytes +/// The `onnx_file` and `tokenizer_files` fields contain the file contents as bytes. #[derive(Debug, Clone, PartialEq, Eq)] pub struct UserDefinedBgem3Model { pub onnx_file: Vec, diff --git a/src/common.rs b/src/common.rs index d9761c6..e2aba9d 100644 --- a/src/common.rs +++ b/src/common.rs @@ -107,8 +107,8 @@ pub struct TokenizerFiles { pub tokenizer_config_file: Vec, } -/// The procedure for loading tokenizer files from the hugging face hub is separated -/// from the main load_tokenizer function (which is expecting bytes, from any source). +/// Loading tokenizer files from the Hugging Face Hub is separate from the main +/// `load_tokenizer` function, which expects bytes from any source. #[cfg(feature = "hf-hub")] pub fn load_tokenizer_hf_hub(model_repo: ApiRepo, max_length: usize) -> Result { let tokenizer_files: TokenizerFiles = TokenizerFiles { @@ -159,7 +159,8 @@ pub fn load_tokenizer(tokenizer_files: TokenizerFiles, max_length: usize) -> Res ) })?; - //For BGEBaseSmall, the model_max_length value is set to 1000000000000000019884624838656. Which fits in a f64 + // For BGEBaseSmall, model_max_length is set to 1000000000000000019884624838656. + // This value fits in an f64. let model_max_length = tokenizer_config["model_max_length"] .as_f64() .ok_or_else(|| { diff --git a/src/image_embedding/init.rs b/src/image_embedding/init.rs index dd18cd2..e958352 100644 --- a/src/image_embedding/init.rs +++ b/src/image_embedding/init.rs @@ -70,7 +70,7 @@ impl From for ImageInitOptionsUserDefined { /// Struct for "bring your own" embedding models /// -/// The onnx_file and preprocessor_files are expecting the files' bytes +/// The `onnx_file` and `preprocessor_file` fields contain the file contents as bytes. #[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] pub struct UserDefinedImageEmbeddingModel { diff --git a/src/image_embedding/utils.rs b/src/image_embedding/utils.rs index 8791f44..4d56a15 100644 --- a/src/image_embedding/utils.rs +++ b/src/image_embedding/utils.rs @@ -144,7 +144,7 @@ impl Transform for CenterCrop { Array3::zeros((3usize, crop_height as usize, crop_width as usize)); let offset_x = (crop_width - origin_width) / 2; let offset_y = (crop_height - origin_height) / 2; - // whc -> chw + // HWC -> CHW for (x, y, pixel) in image.to_rgb8().enumerate_pixels() { pixels_array[[0, (y + offset_y) as usize, (x + offset_x) as usize]] = pixel[0] as f32; @@ -166,7 +166,7 @@ impl Transform for PILToNDarray { TransformData::Image(image) => { let image = image.to_rgb8(); let (width, height) = image.dimensions(); - // whc -> chw + // HWC -> CHW let mut pixels_array = Array3::zeros((3usize, height as usize, width as usize)); for (x, y, pixel) in image.enumerate_pixels() { pixels_array[[0, y as usize, x as usize]] = pixel[0] as f32; diff --git a/src/models/image_embedding.rs b/src/models/image_embedding.rs index 40476ed..ffeae44 100644 --- a/src/models/image_embedding.rs +++ b/src/models/image_embedding.rs @@ -40,7 +40,7 @@ pub fn models_list() -> Vec> { ModelInfo { model: ImageEmbeddingModel::UnicomVitB16, dim: 768, - description: String::from("Unicom Unicom-ViT-B-16 from open-metric-learning"), + description: String::from("Unicom-ViT-B-16 from Open Metric Learning"), model_code: String::from("Qdrant/Unicom-ViT-B-16"), model_file: String::from("model.onnx"), additional_files: Vec::new(), @@ -49,7 +49,7 @@ pub fn models_list() -> Vec> { ModelInfo { model: ImageEmbeddingModel::UnicomVitB32, dim: 512, - description: String::from("Unicom Unicom-ViT-B-32 from open-metric-learning"), + description: String::from("Unicom-ViT-B-32 from Open Metric Learning"), model_code: String::from("Qdrant/Unicom-ViT-B-32"), model_file: String::from("model.onnx"), additional_files: Vec::new(), @@ -58,7 +58,7 @@ pub fn models_list() -> Vec> { ModelInfo { model: ImageEmbeddingModel::NomicEmbedVisionV15, dim: 768, - description: String::from("Nomic NomicEmbedVisionV15"), + description: String::from("Nomic Embed Vision v1.5"), model_code: String::from("nomic-ai/nomic-embed-vision-v1.5"), model_file: String::from("onnx/model.onnx"), additional_files: Vec::new(), diff --git a/src/models/reranking.rs b/src/models/reranking.rs index 9200e3b..10fa19f 100644 --- a/src/models/reranking.rs +++ b/src/models/reranking.rs @@ -12,7 +12,13 @@ pub enum RerankerModel { /// jinaai/jina-reranker-v1-turbo-en JINARerankerV1TurboEn, /// jinaai/jina-reranker-v2-base-multilingual - JINARerankerV2BaseMultiligual, + JINARerankerV2BaseMultilingual, +} + +impl RerankerModel { + #[allow(non_upper_case_globals)] + #[deprecated(note = "use `JINARerankerV2BaseMultilingual` instead")] + pub const JINARerankerV2BaseMultiligual: Self = Self::JINARerankerV2BaseMultilingual; } pub fn reranker_model_list() -> Vec { @@ -26,7 +32,7 @@ pub fn reranker_model_list() -> Vec { }, RerankerModelInfo { model: RerankerModel::BGERerankerV2M3, - description: String::from("reranker model for multilingual"), + description: String::from("Multilingual reranker model"), model_code: String::from("rozgo/bge-reranker-v2-m3"), model_file: String::from("model.onnx"), additional_files: vec![String::from("model.onnx.data")], @@ -39,8 +45,8 @@ pub fn reranker_model_list() -> Vec { additional_files: vec![], }, RerankerModelInfo { - model: RerankerModel::JINARerankerV2BaseMultiligual, - description: String::from("reranker model for multilingual"), + model: RerankerModel::JINARerankerV2BaseMultilingual, + description: String::from("Multilingual reranker model"), model_code: String::from("jinaai/jina-reranker-v2-base-multilingual"), model_file: String::from("onnx/model.onnx"), additional_files: vec![], @@ -86,14 +92,14 @@ pub(crate) fn all_variants() -> Vec { RerankerModel::BGERerankerBase => (), RerankerModel::BGERerankerV2M3 => (), RerankerModel::JINARerankerV1TurboEn => (), - RerankerModel::JINARerankerV2BaseMultiligual => (), + RerankerModel::JINARerankerV2BaseMultilingual => (), } } vec![ RerankerModel::BGERerankerBase, RerankerModel::BGERerankerV2M3, RerankerModel::JINARerankerV1TurboEn, - RerankerModel::JINARerankerV2BaseMultiligual, + RerankerModel::JINARerankerV2BaseMultilingual, ] } diff --git a/src/models/sparse.rs b/src/models/sparse.rs index f4d5638..b335fc2 100644 --- a/src/models/sparse.rs +++ b/src/models/sparse.rs @@ -31,7 +31,7 @@ pub fn models_list() -> Vec> { model: SparseModel::BGEM3, dim: 0, description: String::from( - "BGE-M3 sparse embedding model with 8192 context, supports 100+ languages", + "BGE-M3 sparse embedding model with an 8192-token context, supporting 100+ languages", ), model_code: String::from("BAAI/bge-m3"), model_file: String::from("onnx/model.onnx"), diff --git a/src/reranking/init.rs b/src/reranking/init.rs index 9264c30..64b28b5 100644 --- a/src/reranking/init.rs +++ b/src/reranking/init.rs @@ -21,9 +21,9 @@ impl HasMaxLength for RerankerModel { /// Options for initializing the reranking models pub type RerankInitOptions = InitOptionsWithLength; -/// Options for initializing UserDefinedRerankerModel +/// Options for initializing a user-defined reranking model. /// -/// Model files are held by the UserDefinedRerankerModel struct +/// Model files are held by the [`UserDefinedRerankingModel`] struct. #[derive(Debug, Clone)] #[non_exhaustive] pub struct RerankInitOptionsUserDefined { @@ -144,7 +144,7 @@ impl From for OnnxSource { /// Struct for "bring your own" reranking models /// -/// The onnx_file and tokenizer_files are expecting the files' bytes +/// The `onnx_file` and `tokenizer_files` fields contain the file contents as bytes. #[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] pub struct UserDefinedRerankingModel { @@ -174,7 +174,7 @@ mod tests { use super::*; #[test] - fn userdefined_builders_set_fields() { + fn user_defined_builders_set_fields() { let o = RerankInitOptionsUserDefined::new() .with_max_length(128) .with_intra_threads(2) diff --git a/src/sparse_text_embedding/init.rs b/src/sparse_text_embedding/init.rs index 7ad2c48..6ff5073 100644 --- a/src/sparse_text_embedding/init.rs +++ b/src/sparse_text_embedding/init.rs @@ -19,7 +19,7 @@ pub type SparseInitOptions = InitOptionsWithLength; /// Struct for "bring your own" sparse embedding models /// -/// The onnx_file and tokenizer_files are expecting the files' bytes +/// The `onnx_file` and `tokenizer_files` fields contain the file contents as bytes. #[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] pub struct UserDefinedSparseModel { diff --git a/src/text_embedding/impl.rs b/src/text_embedding/impl.rs index 7701783..cd90047 100644 --- a/src/text_embedding/impl.rs +++ b/src/text_embedding/impl.rs @@ -63,7 +63,7 @@ impl TextEmbedding { })?; } - // prioritise loading pooling config if available, if not (thanks qdrant!), look for it in hardcoded + // Prioritize loading the pooling config if available; otherwise, look for it in the hardcoded defaults. let post_processing = TextEmbedding::get_default_pooling_method(&model_name); let session = init_session_builder(execution_providers, intra_threads, session_config)? diff --git a/src/text_embedding/init.rs b/src/text_embedding/init.rs index 5c6b990..5d5f026 100644 --- a/src/text_embedding/init.rs +++ b/src/text_embedding/init.rs @@ -22,7 +22,7 @@ impl HasMaxLength for EmbeddingModel { /// Options for initializing the TextEmbedding model pub type TextInitOptions = InitOptionsWithLength; -/// Options for initializing UserDefinedEmbeddingModel +/// Options for initializing a user-defined embedding model. /// /// Model files are held by the UserDefinedEmbeddingModel struct #[derive(Debug, Clone)] @@ -150,7 +150,7 @@ impl From for InitOptionsUserDefined { /// Struct for "bring your own" embedding models /// -/// The onnx_file and tokenizer_files are expecting the files' bytes +/// The `onnx_file` and `tokenizer_files` fields contain the file contents as bytes. #[derive(Debug, Clone, PartialEq, Eq)] pub struct UserDefinedEmbeddingModel { pub onnx_file: Vec, @@ -163,8 +163,8 @@ pub struct UserDefinedEmbeddingModel { /// Struct for adding external initializers to "bring your own" embedding models /// -/// The buffer is expecting the data of the external initializer and the file_name -/// must match the one referenced by the model. +/// The buffer contains the external initializer's data, and `file_name` must +/// match the name referenced by the model. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ExternalInitializerFile { pub file_name: String, diff --git a/tests/text-embeddings.rs b/tests/text-embeddings.rs index 49f2384..4083ce7 100644 --- a/tests/text-embeddings.rs +++ b/tests/text-embeddings.rs @@ -38,7 +38,7 @@ fn verify_embeddings(model: &EmbeddingModel, embeddings: &[Embedding]) -> Result EmbeddingModel::AllMiniLML12V2Q => [-0.07808663, 0.27919534, -0.0770612, -0.75660324], EmbeddingModel::AllMiniLML6V2 => [0.59605527, 0.36542925, -0.16450031, -0.40903988], EmbeddingModel::AllMiniLML6V2Q => [0.5677276, 0.40180072, -0.15454668, -0.4672576], - EmbeddingModel::AllMpnetBaseV2=> [-0.21253541, -0.050802127, 0.14072442, -0.2908188], + EmbeddingModel::AllMpnetBaseV2 => [-0.21253541, -0.050802127, 0.14072442, -0.2908188], EmbeddingModel::BGEBaseENV15 => [-0.51290065, -0.4844747, -0.53036124, -0.5337459], EmbeddingModel::BGEBaseENV15Q => [-0.5130697, -0.48461288, -0.53067875, -0.5337806], EmbeddingModel::BGELargeENV15 => [-0.19347441, -0.28394595, -0.1549195, -0.22201893], @@ -547,7 +547,7 @@ fn test_bgesmallen1point5_match_python_counterpart() { // Tokenize with python SentenceTransformer("BAAI/bge-small-en-v1.5") default tokenizer // with (text, padding="max_length",max_length=384,truncation=True, return_tensors="np"). // Normalized and pooled with SentenceTransformer("BAAI/bge-small-en-v1.5") default pooling settings. - // we only take a 10 items to keep the test file polite + // We only take 10 items to keep the test file small. let baseline: Vec = vec![ 4.208_193_7e-2, -2.748_133_2e-2, @@ -586,7 +586,7 @@ fn test_allminilml6v2_match_python_counterpart() { // Tokenizer with python SentenceTransformer("all-mini-lm-l6-v2") default tokenizer // with (text, padding="max_length",max_length=384,truncation=True, return_tensors="np"). // Normalized and pooled with SentenceTransformer("all-mini-lm-l6-v2") default pooling settings. - // we only take a 10 items to keep the test file polite + // We only take 10 items to keep the test file small. let baseline: Vec = vec![ 3.510_517_6e-2, 1.046_043e-2,