docs: Updated README.md

* Opened image support

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>

README.md

@@ -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 [supported](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]
@@ -209,4 +224,4 @@ search_result = client.query(
     query_text="This is a query document"
 )
 print(search_result)
-```
+```

RELEASE.md

@@ -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

docs/examples/Supported_Models.ipynb

@@ -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": {

fastembed/common/__init__.py

@@ -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"]

fastembed/common/model_management.py

@@ -1,4 +1,5 @@
 import os
+import time
 import shutil
 import tarfile
 from pathlib import Path
@@ -42,9 +43,7 @@ class ModelManagement:
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
 
     @classmethod
-    def download_file_from_gcs(
-        cls, url: str, output_path: str, show_progress: bool = True
-    ) -> str:
+    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
         """
         Downloads a file from Google Cloud Storage.
 
@@ -73,9 +72,7 @@ class ModelManagement:
 
         # Warn if the total size is zero
         if total_size_in_bytes == 0:
-            print(
-                f"Warning: Content-length header is missing or zero in the response from {url}."
-            )
+            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
 
         show_progress = total_size_in_bytes and show_progress
 
@@ -163,9 +160,7 @@ class ModelManagement:
         return cache_dir
 
     @classmethod
-    def retrieve_model_gcs(
-        cls, model_name: str, source_url: str, cache_dir: str
-    ) -> Path:
+    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
 
         cache_tmp_dir = Path(cache_dir) / "tmp"
@@ -191,12 +186,8 @@ class ModelManagement:
             output_path=str(model_tar_gz),
         )
 
-        cls.decompress_to_cache(
-            targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir)
-        )
-        assert (
-            model_tmp_dir.exists()
-        ), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
+        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
 
         model_tar_gz.unlink()
         # Rename from tmp to final name is atomic
@@ -205,7 +196,7 @@ class ModelManagement:
         return model_dir
 
     @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path, **kwargs) -> Path:
+    def download_model(cls, model: Dict[str, Any], cache_dir: Path, retries=3, **kwargs) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
@@ -225,6 +216,7 @@ class ModelManagement:
                 }
                 ```
             cache_dir (str): The path to the cache directory.
+            retries: (int): The number of times to retry (including the first attempt)
 
         Returns:
             Path: The path to the downloaded model directory.
@@ -233,26 +225,38 @@ class ModelManagement:
         hf_source = model.get("sources", {}).get("hf")
         url_source = model.get("sources", {}).get("url")
 
-        if hf_source:
-            extra_patterns = [model["model_file"]]
-            extra_patterns.extend(model.get("additional_files", []))
+        sleep = 3.0
+        while retries > 0:
+            retries -= 1
 
-            try:
-                return Path(
-                    cls.download_files_from_huggingface(
-                        hf_source,
-                        cache_dir=str(cache_dir),
-                        extra_patterns=extra_patterns,
-                        local_files_only=kwargs.get("local_files_only", False),
+            if hf_source:
+                extra_patterns = [model["model_file"]]
+                extra_patterns.extend(model.get("additional_files", []))
+
+                try:
+                    return Path(
+                        cls.download_files_from_huggingface(
+                            hf_source,
+                            cache_dir=str(cache_dir),
+                            extra_patterns=extra_patterns,
+                            local_files_only=kwargs.get("local_files_only", False),
+                        )
                     )
-                )
-            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
-                logger.error(
-                    f"Could not download model from HuggingFace: {e}"
-                    "Falling back to other sources."
-                )
+                except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
+                    logger.error(
+                        f"Could not download model from HuggingFace: {e} "
+                        "Falling back to other sources."
+                    )
+            if url_source:
+                try:
+                    return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
+                except Exception:
+                    logger.error(f"Could not download model from url: {url_source}")
 
-        if url_source:
-            return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
+            logger.error(
+                f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
+            )
+            time.sleep(sleep)
+            sleep *= 3
 
         raise ValueError(f"Could not download model {model['model']} from any source.")

fastembed/common/types.py

@@ -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]]]

fastembed/image/image_embedding.py

@@ -65,8 +65,8 @@ class ImageEmbedding(ImageEmbeddingBase):
                 return
 
         raise ValueError(
-            f"Model {model_name} is not supported in TextEmbedding."
-            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+            f"Model {model_name} is not supported in ImageEmbedding."
+            "Please check the supported models using `ImageEmbedding.list_supported_models()`"
         )
 
     def embed(

fastembed/image/onnx_embedding.py

@@ -29,6 +29,26 @@ supported_onnx_models = [
         },
         "model_file": "model.onnx",
     },
+    {
+        "model": "Qdrant/Unicom-ViT-B-16",
+        "dim": 768,
+        "description": "Unicom Unicom-ViT-B-16 from open-metric-learning",
+        "size_in_GB": 0.82,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-16",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-32",
+        "dim": 512,
+        "description": "Unicom Unicom-ViT-B-32 from open-metric-learning",
+        "size_in_GB": 0.48,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-32",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
 

fastembed/image/onnx_image_model.py

@@ -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

fastembed/sparse/bm25.py

@@ -27,6 +27,7 @@ supported_bm25_models = [
         },
         "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
         "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
     },
 ]
 

fastembed/sparse/bm42.py

@@ -27,6 +27,7 @@ supported_bm42_models = [
         },
         "model_file": "model.onnx",
         "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
     },
 ]
 

fastembed/text/mini_lm_embedding.py

@@ -1,58 +0,0 @@
-from typing import Any, Dict, Iterable, List, Type
-
-import numpy as np
-
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize
-from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
-
-supported_mini_lm_models = [
-    {
-        "model": "sentence-transformers/all-MiniLM-L6-v2",
-        "dim": 384,
-        "description": "Sentence Transformer model, MiniLM-L6-v2",
-        "size_in_GB": 0.09,
-        "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
-            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
-        },
-        "model_file": "model.onnx",
-    }
-]
-
-
-class MiniLMOnnxEmbedding(OnnxTextEmbedding):
-    @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return MiniLMEmbeddingWorker
-
-    @classmethod
-    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
-        token_embeddings = model_output
-        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
-        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
-        input_mask_expanded = input_mask_expanded.astype(float)
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        sum_mask = np.sum(input_mask_expanded, axis=1)
-        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
-        return pooled_embeddings
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
-        """Lists the supported models.
-
-        Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
-        """
-        return supported_mini_lm_models
-
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
-        embeddings = output.model_output
-        attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
-
-
-class MiniLMEmbeddingWorker(OnnxTextEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxTextEmbedding:
-        return MiniLMOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

fastembed/text/onnx_embedding.py

@@ -80,36 +80,6 @@ supported_onnx_models = [
         },
         "model_file": "model_optimized.onnx",
     },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1",
-        "dim": 768,
-        "description": "8192 context length english model",
-        "size_in_GB": 0.52,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1",
-        },
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1.5",
-        "dim": 768,
-        "description": "8192 context length english model",
-        "size_in_GB": 0.52,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1.5",
-        },
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
-        "dim": 768,
-        "description": "Quantized 8192 context length english model",
-        "size_in_GB": 0.13,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1.5",
-        },
-        "model_file": "onnx/model_quantized.onnx",
-    },
     {
         "model": "thenlper/gte-large",
         "dim": 1024,
@@ -274,7 +244,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[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]:
         embeddings = output.model_output
         return normalize(embeddings[:, 0]).astype(np.float32)
 
@@ -286,4 +258,6 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
         cache_dir: str,
         **kwargs,
     ) -> OnnxTextEmbedding:
-        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+        return OnnxTextEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed/text/pooled_embedding.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Iterable, List, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_pooled_models = [
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Quantized 8192 context length english model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class PooledEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return PooledEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(
+        cls, model_output: np.ndarray, attention_mask: np.ndarray
+    ) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(
+            input_mask_expanded, (1, 1, token_embeddings.shape[-1])
+        )
+        input_mask_expanded = input_mask_expanded.astype(float)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_pooled_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext
+    ) -> Iterable[np.ndarray]:
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
+
+
+class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs
+    ) -> OnnxTextEmbedding:
+        return PooledEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed/text/pooled_normalized_embedding.py

@@ -6,8 +6,20 @@ from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import normalize
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.text.pooled_embedding import PooledEmbedding
 
-supported_jina_models = [
+supported_pooled_normalized_models = [
+    {
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "description": "Sentence Transformer model, MiniLM-L6-v2",
+        "size_in_GB": 0.09,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
+            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
+        },
+        "model_file": "model.onnx",
+    },
     {
         "model": "jinaai/jina-embeddings-v2-base-en",
         "dim": 768,
@@ -32,23 +44,21 @@ supported_jina_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",
+    },
 ]
 
 
-class JinaOnnxEmbedding(OnnxTextEmbedding):
+class PooledNormalizedEmbedding(PooledEmbedding):
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return JinaEmbeddingWorker
-
-    @classmethod
-    def mean_pooling(cls, model_output, attention_mask) -> np.ndarray:
-        token_embeddings = model_output
-        input_mask_expanded = (np.expand_dims(attention_mask, axis=-1)).astype(float)
-
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        mask_sum = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
-
-        return sum_embeddings / mask_sum
+        return PooledNormalizedEmbeddingWorker
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -57,7 +67,7 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
         Returns:
             List[Dict[str, Any]]: A list of dictionaries containing the model information.
         """
-        return supported_jina_models
+        return supported_pooled_normalized_models
 
     def _post_process_onnx_output(
         self, output: OnnxOutputContext
@@ -67,10 +77,10 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
         return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
 
 
-class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
+class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
         self, model_name: str, cache_dir: str, **kwargs
     ) -> OnnxTextEmbedding:
-        return JinaOnnxEmbedding(
+        return PooledNormalizedEmbedding(
             model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
         )

fastembed/text/text_embedding.py

@@ -5,8 +5,8 @@ import numpy as np
 from fastembed.common import OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
-from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
-from fastembed.text.mini_lm_embedding import MiniLMOnnxEmbedding
+from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
+from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 
@@ -15,9 +15,9 @@ class TextEmbedding(TextEmbeddingBase):
     EMBEDDINGS_REGISTRY: List[Type[TextEmbeddingBase]] = [
         OnnxTextEmbedding,
         E5OnnxEmbedding,
-        JinaOnnxEmbedding,
         CLIPOnnxEmbedding,
-        MiniLMOnnxEmbedding,
+        PooledNormalizedEmbedding,
+        PooledEmbedding,
     ]
 
     @classmethod

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.3.3"
+version = "0.3.4"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -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"

tests/test_image_onnx_embeddings.py

@@ -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
@@ -11,6 +14,12 @@ CANONICAL_VECTOR_VALUES = {
     "Qdrant/resnet50-onnx": np.array(
         [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]
+    ),
+    "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]
+    ),
 }
 
 
@@ -25,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"]]
 
@@ -36,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")])
@@ -56,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)
 
@@ -68,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)

tests/test_text_onnx_embeddings.py

@@ -45,16 +45,19 @@ CANONICAL_VECTOR_VALUES = {
         [-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]
     ),
     "jinaai/jina-embeddings-v2-base-de": np.array(
-      [-0.0085, 0.0417, 0.0342, 0.0309, -0.0149]
+        [-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.0061, 0.0103, -0.0296, -0.0242, -0.0170]
+        [0.3708 ,  0.2031, -0.3406, -0.2114, -0.3230]
     ),
     "nomic-ai/nomic-embed-text-v1.5": np.array(
-        [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
+        [-0.15407836, -0.03053198, -3.9138033, 0.1910364, 0.13224715]
     ),
     "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
-        [-0.01554983, 0.0129992, -0.17909265, -0.01062993, 0.00512859]
+        [-0.12525563,  0.38030425, -3.961622 ,  0.04176439, -0.0758301]
     ),
     "thenlper/gte-large": np.array(
         [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]