mirror of
https://github.com/qdrant/fastembed.git
synced 2026-09-24 15:07:50 -05:00
Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Anush008
|
6fe808628a | ||
|
Anush
|
9a828da000 | ||
|
Anush008
|
63b2dad4d7 | ||
|
Dmitrii OgnandGeorge Panchuk
|
9c72d2f59f | ||
|
generall
|
9d2175e97b | ||
|
Anush
|
0e258ab875 | ||
|
George
|
e49789c129 | ||
|
Anush
|
1bf72922ce | ||
|
Anush
|
70566dff99 |
@@ -6,11 +6,11 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
|
||||
|
||||
## 📈 Why FastEmbed?
|
||||
|
||||
1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda.
|
||||
1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda.
|
||||
|
||||
2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
|
||||
2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data parallelism for encoding large datasets.
|
||||
|
||||
3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
|
||||
3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever-expanding set of models, including a few multilingual models.
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
@@ -48,13 +48,14 @@ len(embeddings_list[0]) # Vector of 384 dimensions
|
||||
|
||||
Fastembed supports a variety of models for different tasks and modalities.
|
||||
The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
|
||||
|
||||
### 🎒 Dense text embeddings
|
||||
|
||||
```python
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
|
||||
embeddings = list(embedding_model.embed(documents))
|
||||
embeddings = list(model.embed(documents))
|
||||
|
||||
# [
|
||||
# array([-0.1115, 0.0097, 0.0052, 0.0195, ...], dtype=float32),
|
||||
@@ -63,8 +64,6 @@ embeddings = list(embedding_model.embed(documents))
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
### 🔱 Sparse text embeddings
|
||||
|
||||
* SPLADE++
|
||||
@@ -73,7 +72,7 @@ embeddings = list(embedding_model.embed(documents))
|
||||
from fastembed import SparseTextEmbedding
|
||||
|
||||
model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
|
||||
embeddings = list(embedding_model.embed(documents))
|
||||
embeddings = list(model.embed(documents))
|
||||
|
||||
# [
|
||||
# SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
|
||||
@@ -81,30 +80,47 @@ embeddings = list(embedding_model.embed(documents))
|
||||
# ]
|
||||
```
|
||||
|
||||
<!--
|
||||
* BM42 - ([link](ToDo))
|
||||
* BM25
|
||||
|
||||
```python
|
||||
from fastembed import SparseTextEmbedding
|
||||
|
||||
model = SparseTextEmbedding(model_name="Qdrant/bm25")
|
||||
embeddings = list(model.embed(documents))
|
||||
|
||||
# [
|
||||
# SparseEmbedding(indices=[ 129793020, 1999429279, 819028769, ... ], values=[1.6477, 1.6327, 1.2377, ...]),
|
||||
# SparseEmbedding(indices=[ 682147660, 1100855371, 339478471, ... ], values=[1.6741, 1.5432, 1.6741, ...])
|
||||
# ]
|
||||
```
|
||||
|
||||
* [BM42](https://qdrant.tech/articles/bm42/)
|
||||
|
||||
```python
|
||||
from fastembed import SparseTextEmbedding
|
||||
|
||||
model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
|
||||
embeddings = list(embedding_model.embed(documents))
|
||||
embeddings = list(model.embed(documents))
|
||||
|
||||
# [
|
||||
# SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
|
||||
# SparseEmbedding(indices=[ 38, 12, 91, ... ], values=[0.11, 0.22, 0.39, ...])
|
||||
# ]
|
||||
```
|
||||
-->
|
||||
|
||||
You can install [PyStemmer](https://pypi.org/project/PyStemmer/) to improve the stemming performance when using BM25, BM42.
|
||||
|
||||
```shell
|
||||
pip install fastembed[pystemmer]
|
||||
```
|
||||
|
||||
### 🦥 Late interaction models (aka ColBERT)
|
||||
|
||||
|
||||
```python
|
||||
from fastembed import LateInteractionTextEmbedding
|
||||
|
||||
model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
|
||||
embeddings = list(embedding_model.embed(documents))
|
||||
embeddings = list(model.embed(documents))
|
||||
|
||||
# [
|
||||
# array([
|
||||
@@ -129,7 +145,7 @@ images = [
|
||||
]
|
||||
|
||||
model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
|
||||
embeddings = list(embedding_model.embed(images))
|
||||
embeddings = list(model.embed(images))
|
||||
|
||||
# [
|
||||
# array([-0.1115, 0.0097, 0.0052, 0.0195, ...], dtype=float32),
|
||||
@@ -137,7 +153,6 @@ embeddings = list(embedding_model.embed(images))
|
||||
# ]
|
||||
```
|
||||
|
||||
|
||||
## ⚡️ FastEmbed on a GPU
|
||||
|
||||
FastEmbed supports running on GPU devices.
|
||||
@@ -147,7 +162,7 @@ It requires installation of the `fastembed-gpu` package.
|
||||
pip install fastembed-gpu
|
||||
```
|
||||
|
||||
Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
|
||||
Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions and CUDA 12.x support.
|
||||
|
||||
```python
|
||||
from fastembed import TextEmbedding
|
||||
@@ -168,7 +183,7 @@ Installation with Qdrant Client in Python:
|
||||
pip install qdrant-client[fastembed]
|
||||
```
|
||||
|
||||
or
|
||||
or
|
||||
|
||||
```bash
|
||||
pip install qdrant-client[fastembed-gpu]
|
||||
|
||||
+1
-1
@@ -12,7 +12,7 @@ This is a guide how to release `fastembed` and `fastembed-gpu` packages.
|
||||
```bash
|
||||
git checkout gpu
|
||||
git rebase main
|
||||
git push origin gpu
|
||||
git push -f origin gpu
|
||||
```
|
||||
|
||||
4. Draft release notes
|
||||
|
||||
@@ -54,7 +54,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T18:13:25.863008Z",
|
||||
@@ -106,16 +106,16 @@
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>2</th>\n",
|
||||
" <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-xs</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
|
||||
" <td>Based on all-MiniLM-L6-v2 model with only 22m ...</td>\n",
|
||||
" <td>0.090</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>3</th>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-xs</td>\n",
|
||||
" <td>sentence-transformers/all-MiniLM-L6-v2</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Based on all-MiniLM-L6-v2 model with only 22m ...</td>\n",
|
||||
" <td>Sentence Transformer model, MiniLM-L6-v2</td>\n",
|
||||
" <td>0.090</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
@@ -127,16 +127,16 @@
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>5</th>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-s</td>\n",
|
||||
" <td>BAAI/bge-small-en</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
|
||||
" <td>Fast English model</td>\n",
|
||||
" <td>0.130</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>6</th>\n",
|
||||
" <td>BAAI/bge-small-en</td>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-s</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Fast English model</td>\n",
|
||||
" <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
|
||||
" <td>0.130</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
@@ -169,83 +169,97 @@
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>11</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-de</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>German embedding model supporting 8192 sequenc...</td>\n",
|
||||
" <td>0.320</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>12</th>\n",
|
||||
" <td>BAAI/bge-base-en</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Base English model</td>\n",
|
||||
" <td>0.420</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>12</th>\n",
|
||||
" <th>13</th>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-m</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Based on intfloat/e5-base-unsupervised model, ...</td>\n",
|
||||
" <td>0.430</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>13</th>\n",
|
||||
" <td>nomic-ai/nomic-embed-text-v1</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>8192 context length english model</td>\n",
|
||||
" <td>0.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>14</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-en</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>English embedding model supporting 8192 sequen...</td>\n",
|
||||
" <td>0.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>15</th>\n",
|
||||
" <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>8192 context length english model</td>\n",
|
||||
" <td>0.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>15</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-en</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>English embedding model supporting 8192 sequen...</td>\n",
|
||||
" <td>0.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>16</th>\n",
|
||||
" <td>nomic-ai/nomic-embed-text-v1</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>8192 context length english model</td>\n",
|
||||
" <td>0.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>17</th>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
|
||||
" <td>0.540</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>17</th>\n",
|
||||
" <th>18</th>\n",
|
||||
" <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>MixedBread Base sentence embedding model, does...</td>\n",
|
||||
" <td>0.640</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>18</th>\n",
|
||||
" <th>19</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-code</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Source code embedding model supporting 8192 se...</td>\n",
|
||||
" <td>0.640</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>20</th>\n",
|
||||
" <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Sentence-transformers model for tasks like clu...</td>\n",
|
||||
" <td>1.000</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>19</th>\n",
|
||||
" <th>21</th>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-l</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Based on intfloat/e5-large-unsupervised, large...</td>\n",
|
||||
" <td>1.020</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>20</th>\n",
|
||||
" <td>BAAI/bge-large-en-v1.5</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Large English model, v1.5</td>\n",
|
||||
" <td>1.200</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>21</th>\n",
|
||||
" <th>22</th>\n",
|
||||
" <td>thenlper/gte-large</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Large general text embeddings model</td>\n",
|
||||
" <td>1.200</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>22</th>\n",
|
||||
" <th>23</th>\n",
|
||||
" <td>BAAI/bge-large-en-v1.5</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Large English model, v1.5</td>\n",
|
||||
" <td>1.200</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>24</th>\n",
|
||||
" <td>intfloat/multilingual-e5-large</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Multilingual model, e5-large. Recommend using ...</td>\n",
|
||||
@@ -259,55 +273,59 @@
|
||||
" model dim \\\n",
|
||||
"0 BAAI/bge-small-en-v1.5 384 \n",
|
||||
"1 BAAI/bge-small-zh-v1.5 512 \n",
|
||||
"2 sentence-transformers/all-MiniLM-L6-v2 384 \n",
|
||||
"3 snowflake/snowflake-arctic-embed-xs 384 \n",
|
||||
"2 snowflake/snowflake-arctic-embed-xs 384 \n",
|
||||
"3 sentence-transformers/all-MiniLM-L6-v2 384 \n",
|
||||
"4 jinaai/jina-embeddings-v2-small-en 512 \n",
|
||||
"5 snowflake/snowflake-arctic-embed-s 384 \n",
|
||||
"6 BAAI/bge-small-en 384 \n",
|
||||
"5 BAAI/bge-small-en 384 \n",
|
||||
"6 snowflake/snowflake-arctic-embed-s 384 \n",
|
||||
"7 nomic-ai/nomic-embed-text-v1.5-Q 768 \n",
|
||||
"8 BAAI/bge-base-en-v1.5 768 \n",
|
||||
"9 sentence-transformers/paraphrase-multilingual-... 384 \n",
|
||||
"10 Qdrant/clip-ViT-B-32-text 512 \n",
|
||||
"11 BAAI/bge-base-en 768 \n",
|
||||
"12 snowflake/snowflake-arctic-embed-m 768 \n",
|
||||
"13 nomic-ai/nomic-embed-text-v1 768 \n",
|
||||
"14 jinaai/jina-embeddings-v2-base-en 768 \n",
|
||||
"15 nomic-ai/nomic-embed-text-v1.5 768 \n",
|
||||
"16 snowflake/snowflake-arctic-embed-m-long 768 \n",
|
||||
"17 mixedbread-ai/mxbai-embed-large-v1 1024 \n",
|
||||
"18 sentence-transformers/paraphrase-multilingual-... 768 \n",
|
||||
"19 snowflake/snowflake-arctic-embed-l 1024 \n",
|
||||
"20 BAAI/bge-large-en-v1.5 1024 \n",
|
||||
"21 thenlper/gte-large 1024 \n",
|
||||
"22 intfloat/multilingual-e5-large 1024 \n",
|
||||
"11 jinaai/jina-embeddings-v2-base-de 768 \n",
|
||||
"12 BAAI/bge-base-en 768 \n",
|
||||
"13 snowflake/snowflake-arctic-embed-m 768 \n",
|
||||
"14 nomic-ai/nomic-embed-text-v1.5 768 \n",
|
||||
"15 jinaai/jina-embeddings-v2-base-en 768 \n",
|
||||
"16 nomic-ai/nomic-embed-text-v1 768 \n",
|
||||
"17 snowflake/snowflake-arctic-embed-m-long 768 \n",
|
||||
"18 mixedbread-ai/mxbai-embed-large-v1 1024 \n",
|
||||
"19 jinaai/jina-embeddings-v2-base-code 768 \n",
|
||||
"20 sentence-transformers/paraphrase-multilingual-... 768 \n",
|
||||
"21 snowflake/snowflake-arctic-embed-l 1024 \n",
|
||||
"22 thenlper/gte-large 1024 \n",
|
||||
"23 BAAI/bge-large-en-v1.5 1024 \n",
|
||||
"24 intfloat/multilingual-e5-large 1024 \n",
|
||||
"\n",
|
||||
" description size_in_GB \n",
|
||||
"0 Fast and Default English model 0.067 \n",
|
||||
"1 Fast and recommended Chinese model 0.090 \n",
|
||||
"2 Sentence Transformer model, MiniLM-L6-v2 0.090 \n",
|
||||
"3 Based on all-MiniLM-L6-v2 model with only 22m ... 0.090 \n",
|
||||
"2 Based on all-MiniLM-L6-v2 model with only 22m ... 0.090 \n",
|
||||
"3 Sentence Transformer model, MiniLM-L6-v2 0.090 \n",
|
||||
"4 English embedding model supporting 8192 sequen... 0.120 \n",
|
||||
"5 Based on infloat/e5-small-unsupervised, does n... 0.130 \n",
|
||||
"6 Fast English model 0.130 \n",
|
||||
"5 Fast English model 0.130 \n",
|
||||
"6 Based on infloat/e5-small-unsupervised, does n... 0.130 \n",
|
||||
"7 Quantized 8192 context length english model 0.130 \n",
|
||||
"8 Base English model, v1.5 0.210 \n",
|
||||
"9 Sentence Transformer model, paraphrase-multili... 0.220 \n",
|
||||
"10 CLIP text encoder 0.250 \n",
|
||||
"11 Base English model 0.420 \n",
|
||||
"12 Based on intfloat/e5-base-unsupervised model, ... 0.430 \n",
|
||||
"13 8192 context length english model 0.520 \n",
|
||||
"14 English embedding model supporting 8192 sequen... 0.520 \n",
|
||||
"15 8192 context length english model 0.520 \n",
|
||||
"16 Based on nomic-ai/nomic-embed-text-v1-unsuperv... 0.540 \n",
|
||||
"17 MixedBread Base sentence embedding model, does... 0.640 \n",
|
||||
"18 Sentence-transformers model for tasks like clu... 1.000 \n",
|
||||
"19 Based on intfloat/e5-large-unsupervised, large... 1.020 \n",
|
||||
"20 Large English model, v1.5 1.200 \n",
|
||||
"21 Large general text embeddings model 1.200 \n",
|
||||
"22 Multilingual model, e5-large. Recommend using ... 2.240 "
|
||||
"11 German embedding model supporting 8192 sequenc... 0.320 \n",
|
||||
"12 Base English model 0.420 \n",
|
||||
"13 Based on intfloat/e5-base-unsupervised model, ... 0.430 \n",
|
||||
"14 8192 context length english model 0.520 \n",
|
||||
"15 English embedding model supporting 8192 sequen... 0.520 \n",
|
||||
"16 8192 context length english model 0.520 \n",
|
||||
"17 Based on nomic-ai/nomic-embed-text-v1-unsuperv... 0.540 \n",
|
||||
"18 MixedBread Base sentence embedding model, does... 0.640 \n",
|
||||
"19 Source code embedding model supporting 8192 se... 0.640 \n",
|
||||
"20 Sentence-transformers model for tasks like clu... 1.000 \n",
|
||||
"21 Based on intfloat/e5-large-unsupervised, large... 1.020 \n",
|
||||
"22 Large general text embeddings model 1.200 \n",
|
||||
"23 Large English model, v1.5 1.200 \n",
|
||||
"24 Multilingual model, e5-large. Recommend using ... 2.240 "
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@@ -331,7 +349,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T18:13:27.124747Z",
|
||||
@@ -364,29 +382,41 @@
|
||||
" <th>vocab_size</th>\n",
|
||||
" <th>description</th>\n",
|
||||
" <th>size_in_GB</th>\n",
|
||||
" <th>requires_idf</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <th>0</th>\n",
|
||||
" <td>Qdrant/bm42-all-minilm-l6-v2-attentions</td>\n",
|
||||
" <td>30522</td>\n",
|
||||
" <td>Light sparse embedding model, which assigns an...</td>\n",
|
||||
" <td>0.090</td>\n",
|
||||
" <td>Qdrant/bm25</td>\n",
|
||||
" <td>NaN</td>\n",
|
||||
" <td>BM25 as sparse embeddings meant to be used wit...</td>\n",
|
||||
" <td>0.010</td>\n",
|
||||
" <td>True</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>1</th>\n",
|
||||
" <td>prithvida/Splade_PP_en_v1</td>\n",
|
||||
" <td>30522</td>\n",
|
||||
" <td>Misspelled version of the model. Retained for ...</td>\n",
|
||||
" <td>0.532</td>\n",
|
||||
" <td>Qdrant/bm42-all-minilm-l6-v2-attentions</td>\n",
|
||||
" <td>30522.0</td>\n",
|
||||
" <td>Light sparse embedding model, which assigns an...</td>\n",
|
||||
" <td>0.090</td>\n",
|
||||
" <td>True</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>2</th>\n",
|
||||
" <td>prithvida/Splade_PP_en_v1</td>\n",
|
||||
" <td>30522.0</td>\n",
|
||||
" <td>Misspelled version of the model. Retained for ...</td>\n",
|
||||
" <td>0.532</td>\n",
|
||||
" <td>NaN</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>3</th>\n",
|
||||
" <td>prithivida/Splade_PP_en_v1</td>\n",
|
||||
" <td>30522</td>\n",
|
||||
" <td>30522.0</td>\n",
|
||||
" <td>Independent Implementation of SPLADE++ Model f...</td>\n",
|
||||
" <td>0.532</td>\n",
|
||||
" <td>NaN</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>\n",
|
||||
@@ -394,17 +424,19 @@
|
||||
],
|
||||
"text/plain": [
|
||||
" model vocab_size \\\n",
|
||||
"0 Qdrant/bm42-all-minilm-l6-v2-attentions 30522 \n",
|
||||
"1 prithvida/Splade_PP_en_v1 30522 \n",
|
||||
"2 prithivida/Splade_PP_en_v1 30522 \n",
|
||||
"0 Qdrant/bm25 NaN \n",
|
||||
"1 Qdrant/bm42-all-minilm-l6-v2-attentions 30522.0 \n",
|
||||
"2 prithvida/Splade_PP_en_v1 30522.0 \n",
|
||||
"3 prithivida/Splade_PP_en_v1 30522.0 \n",
|
||||
"\n",
|
||||
" description size_in_GB \n",
|
||||
"0 Light sparse embedding model, which assigns an... 0.090 \n",
|
||||
"1 Misspelled version of the model. Retained for ... 0.532 \n",
|
||||
"2 Independent Implementation of SPLADE++ Model f... 0.532 "
|
||||
" description size_in_GB requires_idf \n",
|
||||
"0 BM25 as sparse embeddings meant to be used wit... 0.010 True \n",
|
||||
"1 Light sparse embedding model, which assigns an... 0.090 True \n",
|
||||
"2 Misspelled version of the model. Retained for ... 0.532 NaN \n",
|
||||
"3 Independent Implementation of SPLADE++ Model f... 0.532 NaN "
|
||||
]
|
||||
},
|
||||
"execution_count": 8,
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@@ -429,7 +461,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T18:14:34.370252Z",
|
||||
@@ -482,7 +514,7 @@
|
||||
"0 colbert-ir/colbertv2.0 128 Late interaction model 0.44"
|
||||
]
|
||||
},
|
||||
"execution_count": 10,
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@@ -507,7 +539,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T18:14:42.501881Z",
|
||||
@@ -558,6 +590,20 @@
|
||||
" <td>CLIP vision encoder based on ViT-B/32</td>\n",
|
||||
" <td>0.34</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>2</th>\n",
|
||||
" <td>Qdrant/Unicom-ViT-B-32</td>\n",
|
||||
" <td>512</td>\n",
|
||||
" <td>Unicom Unicom-ViT-B-32 from open-metric-learning</td>\n",
|
||||
" <td>0.48</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>3</th>\n",
|
||||
" <td>Qdrant/Unicom-ViT-B-16</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Unicom Unicom-ViT-B-16 from open-metric-learning</td>\n",
|
||||
" <td>0.82</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>\n",
|
||||
"</div>"
|
||||
@@ -566,13 +612,17 @@
|
||||
" model dim \\\n",
|
||||
"0 Qdrant/resnet50-onnx 2048 \n",
|
||||
"1 Qdrant/clip-ViT-B-32-vision 512 \n",
|
||||
"2 Qdrant/Unicom-ViT-B-32 512 \n",
|
||||
"3 Qdrant/Unicom-ViT-B-16 768 \n",
|
||||
"\n",
|
||||
" description size_in_GB \n",
|
||||
"0 ResNet-50 from `Deep Residual Learning for Ima... 0.10 \n",
|
||||
"1 CLIP vision encoder based on ViT-B/32 0.34 "
|
||||
"1 CLIP vision encoder based on ViT-B/32 0.34 \n",
|
||||
"2 Unicom Unicom-ViT-B-32 from open-metric-learning 0.48 \n",
|
||||
"3 Unicom Unicom-ViT-B-16 from open-metric-learning 0.82 "
|
||||
]
|
||||
},
|
||||
"execution_count": 12,
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@@ -602,7 +652,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.4"
|
||||
"version": "3.11.8"
|
||||
},
|
||||
"orig_nbformat": 4,
|
||||
"vscode": {
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from fastembed.common.types import ImageInput, OnnxProvider, PathInput
|
||||
from fastembed.common.types import ImageInput, OnnxProvider, PathInput, PilInput
|
||||
|
||||
__all__ = ["OnnxProvider", "ImageInput", "PathInput"]
|
||||
__all__ = ["OnnxProvider", "ImageInput", "PathInput", "PilInput"]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
from PIL import Image
|
||||
from typing import Any, Dict, Iterable, Tuple, Union
|
||||
|
||||
if sys.version_info >= (3, 10):
|
||||
@@ -9,6 +10,7 @@ else:
|
||||
|
||||
|
||||
PathInput: TypeAlias = Union[str, os.PathLike]
|
||||
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
|
||||
PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
|
||||
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
|
||||
|
||||
OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]
|
||||
|
||||
@@ -51,6 +51,7 @@ supported_onnx_models = [
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -141,16 +142,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
|
||||
|
||||
return onnx_input
|
||||
|
||||
def _post_process_onnx_output(
|
||||
self, output: OnnxOutputContext
|
||||
) -> Iterable[np.ndarray]:
|
||||
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
|
||||
return normalize(output.model_output).astype(np.float32)
|
||||
|
||||
|
||||
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
|
||||
def init_embedding(
|
||||
self, model_name: str, cache_dir: str, **kwargs
|
||||
) -> OnnxImageEmbedding:
|
||||
return OnnxImageEmbedding(
|
||||
model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
|
||||
)
|
||||
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
|
||||
return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
|
||||
|
||||
@@ -7,7 +7,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from fastembed.common import ImageInput, OnnxProvider, PathInput
|
||||
from fastembed.common import ImageInput, OnnxProvider
|
||||
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
|
||||
from fastembed.common.preprocessor_utils import load_preprocessor
|
||||
from fastembed.common.utils import iter_batch
|
||||
@@ -54,9 +54,12 @@ class OnnxImageModel(OnnxModel[T]):
|
||||
def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
|
||||
return {node.name: encoded for node in self.model.get_inputs()}
|
||||
|
||||
def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext:
|
||||
def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext:
|
||||
with contextlib.ExitStack():
|
||||
image_files = [Image.open(image) for image in images]
|
||||
image_files = [
|
||||
Image.open(image) if not isinstance(image, Image.Image) else image
|
||||
for image in images
|
||||
]
|
||||
encoded = self.processor(image_files)
|
||||
onnx_input = self._build_onnx_input(encoded)
|
||||
onnx_input = self._preprocess_onnx_input(onnx_input)
|
||||
@@ -75,7 +78,11 @@ class OnnxImageModel(OnnxModel[T]):
|
||||
) -> Iterable[T]:
|
||||
is_small = False
|
||||
|
||||
if isinstance(images, str) or isinstance(images, Path):
|
||||
if (
|
||||
isinstance(images, str)
|
||||
or isinstance(images, Path)
|
||||
or (isinstance(images, Image.Image))
|
||||
):
|
||||
images = [images]
|
||||
is_small = True
|
||||
|
||||
@@ -90,9 +97,7 @@ class OnnxImageModel(OnnxModel[T]):
|
||||
for batch in iter_batch(images, batch_size):
|
||||
yield from self._post_process_onnx_output(self.onnx_embed(batch))
|
||||
else:
|
||||
start_method = (
|
||||
"forkserver" if "forkserver" in get_all_start_methods() else "spawn"
|
||||
)
|
||||
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
|
||||
params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
|
||||
pool = ParallelWorkerPool(
|
||||
parallel, self._get_worker_class(), start_method=start_method
|
||||
|
||||
@@ -44,6 +44,14 @@ supported_pooled_normalized_models = [
|
||||
"sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
|
||||
"model_file": "onnx/model_fp16.onnx",
|
||||
},
|
||||
{
|
||||
"model": "jinaai/jina-embeddings-v2-base-code",
|
||||
"dim": 768,
|
||||
"description": "Source code embedding model supporting 8192 sequence length",
|
||||
"size_in_GB": 0.64,
|
||||
"sources": {"hf": "jinaai/jina-embeddings-v2-base-code"},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
|
||||
+4
-1
@@ -25,8 +25,11 @@ numpy = [
|
||||
]
|
||||
pillow = "^10.3.0"
|
||||
snowballstemmer = "^2.2.0"
|
||||
PyStemmer = "^2.2.0"
|
||||
mmh3 = "^4.0"
|
||||
PyStemmer = { version = "^2.2.0", optional = true }
|
||||
|
||||
[tool.poetry.extras]
|
||||
pystemmer = ["PyStemmer"]
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = "^7.4.2"
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from fastembed import ImageEmbedding
|
||||
from tests.config import TEST_MISC_DIR
|
||||
@@ -12,12 +15,10 @@ CANONICAL_VECTOR_VALUES = {
|
||||
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01046245, 0.01171397, 0.00705971, 0.0]
|
||||
),
|
||||
"Qdrant/Unicom-ViT-B-16": np.array(
|
||||
[ 0.0170, -0.0361, 0.0125, -0.0428, -0.0232, 0.0232, -0.0602, -0.0333,
|
||||
0.0155, 0.0497]
|
||||
[0.0170, -0.0361, 0.0125, -0.0428, -0.0232, 0.0232, -0.0602, -0.0333, 0.0155, 0.0497]
|
||||
),
|
||||
"Qdrant/Unicom-ViT-B-32": np.array(
|
||||
[0.0418, 0.0550, 0.0003, 0.0253, -0.0185, 0.0016, -0.0368, -0.0402,
|
||||
-0.0891, -0.0186]
|
||||
[0.0418, 0.0550, 0.0003, 0.0253, -0.0185, 0.0016, -0.0368, -0.0402, -0.0891, -0.0186]
|
||||
),
|
||||
}
|
||||
|
||||
@@ -33,10 +34,15 @@ def test_embedding():
|
||||
|
||||
model = ImageEmbedding(model_name=model_desc["model"])
|
||||
|
||||
images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")]
|
||||
images = [
|
||||
TEST_MISC_DIR / "image.jpeg",
|
||||
str(TEST_MISC_DIR / "small_image.jpeg"),
|
||||
Image.open((TEST_MISC_DIR / "small_image.jpeg")),
|
||||
Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)),
|
||||
]
|
||||
embeddings = list(model.embed(images))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
assert embeddings.shape == (2, dim)
|
||||
assert embeddings.shape == (len(images), dim)
|
||||
|
||||
canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
|
||||
|
||||
@@ -44,19 +50,24 @@ def test_embedding():
|
||||
embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
|
||||
), model_desc["model"]
|
||||
|
||||
assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
|
||||
def test_batch_embedding(n_dims, model_name):
|
||||
model = ImageEmbedding(model_name=model_name)
|
||||
n_images = 32
|
||||
images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
|
||||
n_images // 2
|
||||
)
|
||||
test_images = [
|
||||
TEST_MISC_DIR / "image.jpeg",
|
||||
str(TEST_MISC_DIR / "small_image.jpeg"),
|
||||
Image.open(TEST_MISC_DIR / "small_image.jpeg"),
|
||||
]
|
||||
images = test_images * n_images
|
||||
|
||||
embeddings = list(model.embed(images, batch_size=10))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
assert embeddings.shape == (n_images, n_dims)
|
||||
assert embeddings.shape == (len(test_images) * n_images, n_dims)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
|
||||
@@ -64,9 +75,12 @@ def test_parallel_processing(n_dims, model_name):
|
||||
model = ImageEmbedding(model_name=model_name)
|
||||
|
||||
n_images = 32
|
||||
images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * (
|
||||
n_images // 2
|
||||
)
|
||||
test_images = [
|
||||
TEST_MISC_DIR / "image.jpeg",
|
||||
str(TEST_MISC_DIR / "small_image.jpeg"),
|
||||
Image.open(TEST_MISC_DIR / "small_image.jpeg"),
|
||||
]
|
||||
images = test_images * n_images
|
||||
embeddings = list(model.embed(images, batch_size=10, parallel=2))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
@@ -76,6 +90,6 @@ def test_parallel_processing(n_dims, model_name):
|
||||
embeddings_3 = list(model.embed(images, batch_size=10, parallel=0))
|
||||
embeddings_3 = np.stack(embeddings_3, axis=0)
|
||||
|
||||
assert embeddings.shape == (n_images, n_dims)
|
||||
assert embeddings.shape == (n_images * len(test_images), n_dims)
|
||||
assert np.allclose(embeddings, embeddings_2, atol=1e-3)
|
||||
assert np.allclose(embeddings, embeddings_3, atol=1e-3)
|
||||
|
||||
@@ -47,6 +47,9 @@ CANONICAL_VECTOR_VALUES = {
|
||||
"jinaai/jina-embeddings-v2-base-de": np.array(
|
||||
[-0.0085, 0.0417, 0.0342, 0.0309, -0.0149]
|
||||
),
|
||||
"jinaai/jina-embeddings-v2-base-code": np.array(
|
||||
[0.0145, -0.0164, 0.0136, -0.0170, 0.0734]
|
||||
),
|
||||
"nomic-ai/nomic-embed-text-v1": np.array(
|
||||
[0.3708 , 0.2031, -0.3406, -0.2114, -0.3230]
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user