Skip to content

Commit a890a61

Browse files
authored
feat: Added ClipVitB32 model to support text-to-image search. (Anush008#126)
1 parent f19a540 commit a890a61

File tree

3 files changed

+13
-0
lines changed

3 files changed

+13
-0
lines changed

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ The default model is Flag Embedding, which is top of the [MTEB](https://huggingf
2828
- [**BAAI/bge-small-en-v1.5**](https://huggingface.co/BAAI/bge-small-en-v1.5) - Default
2929
- [**sentence-transformers/all-MiniLM-L6-v2**](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
3030
- [**mixedbread-ai/mxbai-embed-large-v1**](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1)
31+
- [**Qdrant/clip-ViT-B-32-text**](https://huggingface.co/Qdrant/clip-ViT-B-32-text) - pairs with the image model clip-ViT-B-32-vision for image-to-text search
3132

3233
<details>
3334
<summary>Click to see full List</summary>

src/models/text_embedding.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ pub enum EmbeddingModel {
6363
GTELargeENV15,
6464
/// Quantized Alibaba-NLP/gte-large-en-v1.5
6565
GTELargeENV15Q,
66+
/// Qdrant/clip-ViT-B-32-text
67+
ClipVitB32,
6668
}
6769

6870
/// Centralized function to initialize the models map.
@@ -256,6 +258,13 @@ fn init_models_map() -> HashMap<EmbeddingModel, ModelInfo<EmbeddingModel>> {
256258
model_code: String::from("Alibaba-NLP/gte-large-en-v1.5"),
257259
model_file: String::from("onnx/model_quantized.onnx"),
258260
},
261+
ModelInfo {
262+
model: EmbeddingModel::ClipVitB32,
263+
dim: 512,
264+
description: String::from("CLIP text encoder based on ViT-B/32"),
265+
model_code: String::from("Qdrant/clip-ViT-B-32-text"),
266+
model_file: String::from("model.onnx"),
267+
},
259268
];
260269

261270
// TODO: Use when out in stable
@@ -327,6 +336,8 @@ impl EmbeddingModel {
327336
EmbeddingModel::GTEBaseENV15Q => Some(Pooling::Cls),
328337
EmbeddingModel::GTELargeENV15 => Some(Pooling::Cls),
329338
EmbeddingModel::GTELargeENV15Q => Some(Pooling::Cls),
339+
340+
EmbeddingModel::ClipVitB32 => Some(Pooling::Mean),
330341
}
331342
}
332343

tests/embeddings.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ fn verify_embeddings(model: &EmbeddingModel, embeddings: &[Embedding]) -> Result
6161
EmbeddingModel::ParaphraseMLMiniLML12V2 => [-0.07795018, -0.059113946, -0.043668486, -0.1880083],
6262
EmbeddingModel::ParaphraseMLMiniLML12V2Q => [-0.07749095, -0.058981877, -0.043487836, -0.18775631],
6363
EmbeddingModel::ParaphraseMLMpnetBaseV2 => [0.39132136, 0.49490625, 0.65497226, 0.34237382],
64+
EmbeddingModel::ClipVitB32 => [0.7057363, 1.3549932, 0.46823958, 0.52351093],
6465
_ => panic!("Model {model} not found. If you have just inserted this `EmbeddingModel` variant, please update the expected embeddings."),
6566
};
6667

0 commit comments

Comments
 (0)