Compare commits

..
15 Commits
Author SHA1 Message Date
Arun 432da42c11 fix links (#215) 2024-04-27 22:24:29 +05:30
George 5cde2898bc new: remove slurm environment variables (#213) 2024-04-26 21:45:02 +02:00
AnushandGeorge ab7a99a748 feat: Quantized models (#201)
* feat: Quantized models

* refactor: use model_file for GCS

* refactoring: refactor model downloading (#209)

* refactoring: refactor model downloading

* refactor: update docstring

Co-authored-by: Anush <anushshetty90@gmail.com>

* Update fastembed/common/model_management.py

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

* fix: model_file for Snowflake models

---------

Co-authored-by: George <george.panchuk@qdrant.tech>
2024-04-26 17:27:16 +02:00
Anush 466886a317 ci: Schedule python-tests.yml (#211)
* ci: Schedule python-tests.yml

* ci: use emojis

* ci: Bump action versions python-tests.yml

* ci: python-tests.yml
2024-04-26 10:28:37 +05:30
Anush cc4112d859 feat: Snowflake models (#207)
* feat: Snowflake models

* Added snowflake/snowflake-arctic-embed-m

* docs: snowflake/snowflake-arctic-embed-m
2024-04-19 19:00:10 +05:30
Nirant 864217a7d9 Re-run of identical hardware and generate graphs (#205) 2024-04-18 11:13:18 +05:30
Andrew Green 9bad44368e Workaround for running on SLURM (#198)
* Workaround for running on SLURM

onnxruntime would usually get the number of threads from OMP_NUM_THREADS, but that isn't set on SLURM which handles the number of threads differently.

This addition tries to figure out if we're running under SLURM, and if so sets the session options accordingly using SLURM environment variables instead.

Tested working with latest versions of onnxruntime and fastembed on slurm 23.02.07

* onnxruntime requires number of threads be an integer

Caused by me having mis-matched version from my machine and the slurm cluster :(

* `os.getenv` returns None for unset environment variables, fix logic

Instead of empty string, which I thought it did
2024-04-15 19:14:11 +05:30
dependabot[bot]andNirant a5cab7a31a build(deps): bump idna from 3.6 to 3.7 (#195)
Bumps [idna](https://github.com/kjd/idna) from 3.6 to 3.7.
- [Release notes](https://github.com/kjd/idna/releases)
- [Changelog](https://github.com/kjd/idna/blob/master/HISTORY.rst)
- [Commits](https://github.com/kjd/idna/compare/v3.6...v3.7)

---
updated-dependencies:
- dependency-name: idna
  dependency-type: indirect
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Nirant <NirantK@users.noreply.github.com>
2024-04-15 06:04:54 +05:30
Anush e55c145924 chore: Exclude unused model repo files (#196)
* chore: Exclude unused model files

* fix: blob pattern
2024-04-12 20:15:02 +05:30
dependabot[bot]andNirant ad02f60eea build(deps-dev): bump pillow from 10.2.0 to 10.3.0 (#186)
Bumps [pillow](https://github.com/python-pillow/Pillow) from 10.2.0 to 10.3.0.
- [Release notes](https://github.com/python-pillow/Pillow/releases)
- [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst)
- [Commits](https://github.com/python-pillow/Pillow/compare/10.2.0...10.3.0)

---
updated-dependencies:
- dependency-name: pillow
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Nirant <NirantK@users.noreply.github.com>
2024-04-08 16:21:32 +05:30
NirantandGeorge 2fcec07f1b Add Python version to the ISSUE_TEMPLATE (#188)
* Add Python version

* Update .github/ISSUE_TEMPLATE/bug-report.yml

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

---------

Co-authored-by: George <george.panchuk@qdrant.tech>
2024-04-08 16:05:29 +05:30
George f340a73a3e refactoring: update binary quantization notebook (#180) 2024-04-02 14:45:27 +05:30
Nirant 335f673f3a Update bug-report.yml 2024-04-02 13:51:51 +05:30
Nirant ee7ba0a536 Update bug-report.yml 2024-04-02 13:51:27 +05:30
George 912e95e5c8 fix: remove model archive if model extraction was not finished correctly (#179) 2024-04-01 20:11:32 +02:00
14 changed files with 613 additions and 366 deletions
+12 -2
View File
@@ -15,13 +15,23 @@ body:
value: "A bug happened!"
validations:
required: true
- type: textarea
id: Python version
attributes:
label: What Python version are you on? e.g. python --version
description: Also tell us, what package manager are you using e.g. conda, pip, poetry?
placeholder: Python3.10
validations:
required: true
- type: dropdown
id: version
attributes:
label: Version
description: What version of FastEmbed are you running? python -c "import fastembed; print(fastembed.__version__)". If you're not on the latest, please upgrade and see if the problem persists.
options:
- 0.2.4 (Latest)
- 0.2.6 (Latest)
- 0.2.5
- 0.2.4
- 0.2.3
- 0.2.2
- 0.2.1
@@ -43,4 +53,4 @@ body:
attributes:
label: Relevant stack traces and/or logs
description: Please copy and paste any relevant raised exceptions. This will be automatically formatted into code, so no need for backticks.
render: shell
render: shell
+15 -7
View File
@@ -3,6 +3,8 @@ name: Tests
on:
push:
branches: [ master, main ]
schedule:
- cron: 0 0 * * *
pull_request:
env:
@@ -29,9 +31,9 @@ jobs:
name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
steps:
- uses: actions/checkout@v2
- uses: actions/checkout@v3
- name: Set up Python
uses: actions/setup-python@v2
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
@@ -39,8 +41,14 @@ jobs:
python -m pip install poetry
poetry config virtualenvs.create false
poetry install --no-interaction --no-ansi --without docs
- name: Run tests
run: |
export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
pytest
shell: bash
- name: Install Test Dependencies
run: pip install pytest pytest-md pytest-emoji
- name: Run pytest
uses: pavelzw/pytest-action@v2
with:
verbose: true
emoji: true
job-summary: true
report-title: 'FastEmbed Test Report'
+1 -1
View File
@@ -2,7 +2,7 @@
FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).
## 📈 Why FastEmbed?
File diff suppressed because one or more lines are too long
+97 -52
View File
@@ -17,7 +17,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -35,7 +35,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 2,
"metadata": {},
"outputs": [
{
@@ -89,90 +89,125 @@
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>snowflake/snowflake-arctic-embed-xs</td>\n",
" <td>384</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>4</th>\n",
" <td>jinaai/jina-embeddings-v2-small-en</td>\n",
" <td>512</td>\n",
" <td>English embedding model supporting 8192 sequen...</td>\n",
" <td>0.120</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <th>5</th>\n",
" <td>snowflake/snowflake-arctic-embed-s</td>\n",
" <td>384</td>\n",
" <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
" <td>0.130</td>\n",
" </tr>\n",
" <tr>\n",
" <th>6</th>\n",
" <td>BAAI/bge-small-en</td>\n",
" <td>384</td>\n",
" <td>Fast English model</td>\n",
" <td>0.130</td>\n",
" </tr>\n",
" <tr>\n",
" <th>5</th>\n",
" <th>7</th>\n",
" <td>BAAI/bge-base-en-v1.5</td>\n",
" <td>768</td>\n",
" <td>Base English model, v1.5</td>\n",
" <td>0.210</td>\n",
" </tr>\n",
" <tr>\n",
" <th>6</th>\n",
" <th>8</th>\n",
" <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
" <td>384</td>\n",
" <td>Sentence Transformer model, paraphrase-multili...</td>\n",
" <td>0.220</td>\n",
" </tr>\n",
" <tr>\n",
" <th>7</th>\n",
" <th>9</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>8</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>9</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>10</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>11</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>11</th>\n",
" <th>12</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>13</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>14</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>15</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>12</th>\n",
" <th>16</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>13</th>\n",
" <th>17</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>18</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>14</th>\n",
" <th>19</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>15</th>\n",
" <th>20</th>\n",
" <td>intfloat/multilingual-e5-large</td>\n",
" <td>1024</td>\n",
" <td>Multilingual model, e5-large. Recommend using ...</td>\n",
@@ -187,40 +222,50 @@
"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 jinaai/jina-embeddings-v2-small-en 512 \n",
"4 BAAI/bge-small-en 384 \n",
"5 BAAI/bge-base-en-v1.5 768 \n",
"6 sentence-transformers/paraphrase-multilingual-... 384 \n",
"7 BAAI/bge-base-en 768 \n",
"8 nomic-ai/nomic-embed-text-v1 768 \n",
"9 nomic-ai/nomic-embed-text-v1.5 768 \n",
"10 jinaai/jina-embeddings-v2-base-en 768 \n",
"11 mixedbread-ai/mxbai-embed-large-v1 1024 \n",
"12 sentence-transformers/paraphrase-multilingual-... 768 \n",
"13 BAAI/bge-large-en-v1.5 1024 \n",
"14 thenlper/gte-large 1024 \n",
"15 intfloat/multilingual-e5-large 1024 \n",
"3 snowflake/snowflake-arctic-embed-xs 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",
"7 BAAI/bge-base-en-v1.5 768 \n",
"8 sentence-transformers/paraphrase-multilingual-... 384 \n",
"9 BAAI/bge-base-en 768 \n",
"10 snowflake/snowflake-arctic-embed-m 768 \n",
"11 jinaai/jina-embeddings-v2-base-en 768 \n",
"12 nomic-ai/nomic-embed-text-v1 768 \n",
"13 nomic-ai/nomic-embed-text-v1.5 768 \n",
"14 snowflake/snowflake-arctic-embed-m-long 768 \n",
"15 mixedbread-ai/mxbai-embed-large-v1 1024 \n",
"16 sentence-transformers/paraphrase-multilingual-... 768 \n",
"17 snowflake/snowflake-arctic-embed-l 1024 \n",
"18 BAAI/bge-large-en-v1.5 1024 \n",
"19 thenlper/gte-large 1024 \n",
"20 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 English embedding model supporting 8192 sequen... 0.120 \n",
"4 Fast English model 0.130 \n",
"5 Base English model, v1.5 0.210 \n",
"6 Sentence Transformer model, paraphrase-multili... 0.220 \n",
"7 Base English model 0.420 \n",
"8 8192 context length english model 0.520 \n",
"9 8192 context length english model 0.520 \n",
"10 English embedding model supporting 8192 sequen... 0.520 \n",
"11 MixedBread Base sentence embedding model, does... 0.640 \n",
"12 Sentence-transformers model for tasks like clu... 1.000 \n",
"13 Large English model, v1.5 1.200 \n",
"14 Large general text embeddings model 1.200 \n",
"15 Multilingual model, e5-large. Recommend using ... 2.240 "
"3 Based on all-MiniLM-L6-v2 model with only 22m ... 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",
"7 Base English model, v1.5 0.210 \n",
"8 Sentence Transformer model, paraphrase-multili... 0.220 \n",
"9 Base English model 0.420 \n",
"10 Based on intfloat/e5-base-unsupervised model, ... 0.430 \n",
"11 English embedding model supporting 8192 sequen... 0.520 \n",
"12 8192 context length english model 0.520 \n",
"13 8192 context length english model 0.520 \n",
"14 Based on nomic-ai/nomic-embed-text-v1-unsuperv... 0.540 \n",
"15 MixedBread Base sentence embedding model, does... 0.640 \n",
"16 Sentence-transformers model for tasks like clu... 1.000 \n",
"17 Based on intfloat/e5-large-unsupervised, large... 1.020 \n",
"18 Large English model, v1.5 1.200 \n",
"19 Large general text embeddings model 1.200 \n",
"20 Multilingual model, e5-large. Recommend using ... 2.240 "
]
},
"execution_count": 6,
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
+197 -100
View File
@@ -33,8 +33,13 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"execution_count": 1,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:03.013948Z",
"start_time": "2024-04-01T16:33:01.019043Z"
}
},
"outputs": [],
"source": [
"!pip install qdrant-client pandas dataset --quiet --upgrade"
@@ -42,12 +47,34 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"execution_count": 2,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:03.914729Z",
"start_time": "2024-04-01T16:33:03.015394Z"
}
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/joein/work/qdrant/fastembed/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
}
],
"source": [
"import os\n",
"import random\n",
"import time\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
"from qdrant_client import QdrantClient, models"
"from qdrant_client import QdrantClient, models\n",
"\n",
"random.seed(37)\n",
"np.random.seed(37)"
]
},
{
@@ -61,41 +88,59 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"execution_count": 3,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:09.085853Z",
"start_time": "2024-04-01T16:33:03.912688Z"
}
},
"outputs": [
{
"data": {
"text/plain": "100000"
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import datasets\n",
"\n",
"dataset = datasets.load_dataset(\n",
" \"Qdrant/dbpedia-entities-openai3-text-embedding-3-small-1536-100K\", split=\"train\"\n",
")"
")\n",
"len(dataset)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"len(dataset)\n",
"# dataset[0]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"execution_count": 4,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:33:09.176212Z",
"start_time": "2024-04-01T16:33:09.084550Z"
}
},
"outputs": [
{
"data": {
"text/plain": "True"
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"client = QdrantClient(\n",
" timeout=600,\n",
" prefer_grpc=True,\n",
")\n",
"\n",
"collection_name = \"binary-quantization\"\n",
"client.recreate_collection(\n",
" collection_name=f\"{collection_name}\",\n",
" collection_name=collection_name,\n",
" vectors_config=models.VectorParams(\n",
" size=1536,\n",
" distance=models.Distance.DOT,\n",
@@ -109,29 +154,48 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"execution_count": 5,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.383986Z",
"start_time": "2024-04-01T16:33:09.175725Z"
}
},
"outputs": [],
"source": [
"import os\n",
"def iter_dataset(dataset):\n",
" for point in dataset:\n",
" yield point[\"openai\"], {\"text\": point[\"text\"]}\n",
"\n",
"\n",
"bs = 1000\n",
"for i in range(0, len(dataset), bs):\n",
" client.upload_collection(\n",
" collection_name=collection_name,\n",
" ids=range(i, i + bs),\n",
" vectors=dataset[i : i + bs][\"openai\"],\n",
" payload=[{\"text\": x} for x in dataset[i : i + bs][\"text\"]],\n",
" parallel=max(1, (os.cpu_count() // 2)),\n",
" )"
"vectors, payload = zip(*iter_dataset(dataset))\n",
"client.upload_collection(\n",
" collection_name=collection_name,\n",
" vectors=vectors,\n",
" payload=payload,\n",
" parallel=max(1, (os.cpu_count() // 2)),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"execution_count": 6,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.390886Z",
"start_time": "2024-04-01T16:34:13.385961Z"
}
},
"outputs": [
{
"data": {
"text/plain": "{'status': <CollectionStatus.YELLOW: 'yellow'>,\n 'optimizer_status': <OptimizersStatusOneOf.OK: 'ok'>,\n 'vectors_count': 116640,\n 'indexed_vectors_count': 43520,\n 'points_count': 116640,\n 'segments_count': 6,\n 'config': {'params': {'vectors': {'size': 1536,\n 'distance': <Distance.DOT: 'Dot'>,\n 'hnsw_config': None,\n 'quantization_config': None,\n 'on_disk': True},\n 'shard_number': 1,\n 'sharding_method': None,\n 'replication_factor': 1,\n 'write_consistency_factor': 1,\n 'read_fan_out_factor': None,\n 'on_disk_payload': True,\n 'sparse_vectors': None},\n 'hnsw_config': {'m': 16,\n 'ef_construct': 100,\n 'full_scan_threshold': 10000,\n 'max_indexing_threads': 0,\n 'on_disk': False,\n 'payload_m': None},\n 'optimizer_config': {'deleted_threshold': 0.2,\n 'vacuum_min_vector_number': 1000,\n 'default_segment_number': 0,\n 'max_segment_size': None,\n 'memmap_threshold': None,\n 'indexing_threshold': 20000,\n 'flush_interval_sec': 5,\n 'max_optimization_threads': None},\n 'wal_config': {'wal_capacity_mb': 32, 'wal_segments_ahead': 0},\n 'quantization_config': {'binary': {'always_ram': True}}},\n 'payload_schema': {}}"
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"collection_info = client.get_collection(collection_name=f\"{collection_name}\")\n",
"collection_info.dict()"
@@ -150,30 +214,41 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"execution_count": 7,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.453626Z",
"start_time": "2024-04-01T16:34:13.391567Z"
}
},
"outputs": [
{
"data": {
"text/plain": "[89391,\n 79659,\n 12006,\n 80978,\n 87219,\n 97885,\n 83155,\n 67504,\n 4645,\n 82711,\n 48395,\n 57375,\n 69208,\n 14136,\n 89515,\n 59880,\n 78730,\n 36952,\n 49620,\n 96486,\n 55473,\n 58179,\n 18926,\n 6489,\n 11931,\n 54146,\n 9850,\n 71259,\n 37825,\n 47331,\n 84964,\n 92399,\n 56669,\n 77042,\n 73744,\n 47993,\n 83780,\n 92429,\n 75114,\n 4463,\n 69030,\n 81185,\n 27950,\n 66217,\n 54652,\n 8260,\n 1151,\n 993,\n 85954,\n 66863,\n 47303,\n 8992,\n 92688,\n 76030,\n 29472,\n 3077,\n 42454,\n 46120,\n 69140,\n 20877,\n 2844,\n 95423,\n 1770,\n 28568,\n 96448,\n 94227,\n 40837,\n 91684,\n 29785,\n 66936,\n 85121,\n 39546,\n 81910,\n 5514,\n 37068,\n 35731,\n 93990,\n 26685,\n 63076,\n 18762,\n 27922,\n 34916,\n 80976,\n 83189,\n 6328,\n 57508,\n 58860,\n 13758,\n 72976,\n 85030,\n 332,\n 34963,\n 85009,\n 31344,\n 11560,\n 58108,\n 85163,\n 17064,\n 44712,\n 45962]"
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import random\n",
"from random import randint\n",
"\n",
"random.seed(37)\n",
"\n",
"query_indices = [randint(0, len(dataset)) for _ in range(100)]\n",
"query_indices = random.sample(range(len(dataset)), 100)\n",
"query_dataset = dataset[query_indices]\n",
"query_indices"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"execution_count": 8,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.453928Z",
"start_time": "2024-04-01T16:34:13.452405Z"
}
},
"outputs": [],
"source": [
"## Add Gaussian noise to any vector\n",
"import numpy as np\n",
"\n",
"np.random.seed(37)\n",
"\n",
"\n",
"def add_noise(vector, noise=0.05):\n",
@@ -182,76 +257,98 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"execution_count": 9,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:34:13.457839Z",
"start_time": "2024-04-01T16:34:13.455431Z"
}
},
"outputs": [],
"source": [
"import time\n",
"\n",
"\n",
"def correct(results, text):\n",
" result_texts = [x.payload[\"text\"] for x in results]\n",
" return text in result_texts\n",
" return text in [x.payload[\"text\"] for x in results]\n",
"\n",
"\n",
"def count_correct(query_dataset, limit=1, oversampling=1, rescore=False):\n",
" correct_results = 0\n",
" for qv, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
" for query_vector, text in zip(query_dataset[\"openai\"], query_dataset[\"text\"]):\n",
" results = client.search(\n",
" collection_name=collection_name,\n",
" query_vector=add_noise(np.array(qv)),\n",
" query_vector=add_noise(np.array(query_vector)),\n",
" limit=limit,\n",
" search_params=models.SearchParams(\n",
" quantization=models.QuantizationSearchParams(\n",
" ignore=False,\n",
" rescore=rescore,\n",
" oversampling=oversampling,\n",
" )\n",
" ),\n",
" )\n",
" correct_results += correct(results, text)\n",
" return correct_results\n",
"\n",
"\n",
"limit_grid = [1, 3, 5, 10, 20, 50]\n",
"# limit_grid = [1, 3, 5]\n",
"oversampling_grid = [1.0, 1.5, 2.0, 3.0, 5.0]\n",
"# oversampling_grid = [1.0, 1.5, 2.0]\n",
"rescore_grid = [False, True]\n",
"results = []\n",
"for limit in limit_grid:\n",
" for oversampling in oversampling_grid:\n",
" for rescore in rescore_grid:\n",
" # print(f\"limit={limit}, oversampling={oversampling}, rescore={rescore}\")\n",
" start = time.time()\n",
" correct_results = count_correct(\n",
" query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
" )\n",
" end = time.time()\n",
" results.append(\n",
" {\n",
" \"limit\": limit,\n",
" \"oversampling\": oversampling,\n",
" \"rescore\": rescore,\n",
" \"correct\": correct_results,\n",
" \"total queries\": len(query_dataset[\"text\"]),\n",
" \"time\": end - start,\n",
" }\n",
" )\n",
"\n",
"results_df = pd.DataFrame(results)\n",
"results_df"
" return correct_results"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"execution_count": 18,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:40:48.347002Z",
"start_time": "2024-04-01T16:40:42.228551Z"
},
"collapsed": false
},
"outputs": [],
"source": [
"df = results_df.copy()\n",
"df[\"candidates\"] = df[\"oversampling\"] * df[\"limit\"]\n",
"df[[\"candidates\", \"rescore\", \"time\"]]\n",
"limit_grid = [1, 3, 10, 20, 50]\n",
"oversampling_grid = [1.0, 3.0, 5.0]\n",
"rescore_grid = [False, True]\n",
"results = []\n",
"\n",
"for limit in limit_grid:\n",
" for oversampling in oversampling_grid:\n",
" for rescore in rescore_grid:\n",
" start = time.perf_counter()\n",
" correct_results = count_correct(\n",
" query_dataset, limit=limit, oversampling=oversampling, rescore=rescore\n",
" )\n",
" end = time.perf_counter()\n",
" results.append(\n",
" {\n",
" \"limit\": limit,\n",
" \"oversampling\": oversampling,\n",
" \"candidates\": int(oversampling * limit),\n",
" \"rescore\": rescore,\n",
" \"accuracy\": correct_results / 100,\n",
" \"total queries\": len(query_dataset[\"text\"]),\n",
" \"time\": end - start,\n",
" }\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {
"ExecuteTime": {
"end_time": "2024-04-01T16:41:55.445405Z",
"start_time": "2024-04-01T16:41:55.442687Z"
}
},
"outputs": [
{
"data": {
"text/html": "<div>\n<style scoped>\n .dataframe tbody tr th:only-of-type {\n vertical-align: middle;\n }\n\n .dataframe tbody tr th {\n vertical-align: top;\n }\n\n .dataframe thead th {\n text-align: right;\n }\n</style>\n<table border=\"1\" class=\"dataframe\">\n <thead>\n <tr style=\"text-align: right;\">\n <th></th>\n <th>candidates</th>\n <th>rescore</th>\n <th>accuracy</th>\n <th>time</th>\n </tr>\n </thead>\n <tbody>\n <tr>\n <th>0</th>\n <td>1</td>\n <td>False</td>\n <td>0.90</td>\n <td>0.221826</td>\n </tr>\n <tr>\n <th>1</th>\n <td>1</td>\n <td>True</td>\n <td>0.91</td>\n <td>0.134167</td>\n </tr>\n <tr>\n <th>2</th>\n <td>3</td>\n <td>False</td>\n <td>0.88</td>\n <td>0.115299</td>\n </tr>\n <tr>\n <th>3</th>\n <td>3</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.209320</td>\n </tr>\n <tr>\n <th>4</th>\n <td>5</td>\n <td>False</td>\n <td>0.84</td>\n <td>0.154485</td>\n </tr>\n <tr>\n <th>5</th>\n <td>5</td>\n <td>True</td>\n <td>0.91</td>\n <td>0.124424</td>\n </tr>\n <tr>\n <th>6</th>\n <td>3</td>\n <td>False</td>\n <td>0.99</td>\n <td>0.121695</td>\n </tr>\n <tr>\n <th>7</th>\n <td>3</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.123257</td>\n </tr>\n <tr>\n <th>8</th>\n <td>9</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.119629</td>\n </tr>\n <tr>\n <th>9</th>\n <td>9</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.119372</td>\n </tr>\n <tr>\n <th>10</th>\n <td>15</td>\n <td>False</td>\n <td>0.90</td>\n <td>0.121621</td>\n </tr>\n <tr>\n <th>11</th>\n <td>15</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.125466</td>\n </tr>\n <tr>\n <th>12</th>\n <td>10</td>\n <td>False</td>\n <td>0.93</td>\n <td>0.135910</td>\n </tr>\n <tr>\n <th>13</th>\n <td>10</td>\n <td>True</td>\n <td>0.95</td>\n <td>0.138135</td>\n </tr>\n <tr>\n <th>14</th>\n <td>30</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.177928</td>\n </tr>\n <tr>\n <th>15</th>\n <td>30</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.254588</td>\n </tr>\n <tr>\n <th>16</th>\n <td>50</td>\n <td>False</td>\n <td>0.94</td>\n <td>0.268659</td>\n </tr>\n <tr>\n <th>17</th>\n <td>50</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.269792</td>\n </tr>\n <tr>\n <th>18</th>\n <td>20</td>\n <td>False</td>\n <td>0.96</td>\n <td>0.249941</td>\n </tr>\n <tr>\n <th>19</th>\n <td>20</td>\n <td>True</td>\n <td>0.96</td>\n <td>0.247138</td>\n </tr>\n <tr>\n <th>20</th>\n <td>60</td>\n <td>False</td>\n <td>0.97</td>\n <td>0.251301</td>\n </tr>\n <tr>\n <th>21</th>\n <td>60</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.256504</td>\n </tr>\n <tr>\n <th>22</th>\n <td>100</td>\n <td>False</td>\n <td>0.98</td>\n <td>0.270049</td>\n </tr>\n <tr>\n <th>23</th>\n <td>100</td>\n <td>True</td>\n <td>0.97</td>\n <td>0.248972</td>\n </tr>\n <tr>\n <th>24</th>\n <td>50</td>\n <td>False</td>\n <td>0.97</td>\n <td>0.306356</td>\n </tr>\n <tr>\n <th>25</th>\n <td>50</td>\n <td>True</td>\n <td>0.98</td>\n <td>0.257544</td>\n </tr>\n <tr>\n <th>26</th>\n <td>150</td>\n <td>False</td>\n <td>0.98</td>\n <td>0.238811</td>\n </tr>\n <tr>\n <th>27</th>\n <td>150</td>\n <td>True</td>\n <td>0.99</td>\n <td>0.263939</td>\n </tr>\n <tr>\n <th>28</th>\n <td>250</td>\n <td>False</td>\n <td>0.99</td>\n <td>0.256558</td>\n </tr>\n <tr>\n <th>29</th>\n <td>250</td>\n <td>True</td>\n <td>1.00</td>\n <td>0.335823</td>\n </tr>\n </tbody>\n</table>\n</div>",
"text/plain": " candidates rescore accuracy time\n0 1 False 0.90 0.221826\n1 1 True 0.91 0.134167\n2 3 False 0.88 0.115299\n3 3 True 0.97 0.209320\n4 5 False 0.84 0.154485\n5 5 True 0.91 0.124424\n6 3 False 0.99 0.121695\n7 3 True 0.96 0.123257\n8 9 False 0.94 0.119629\n9 9 True 0.98 0.119372\n10 15 False 0.90 0.121621\n11 15 True 0.97 0.125466\n12 10 False 0.93 0.135910\n13 10 True 0.95 0.138135\n14 30 False 0.94 0.177928\n15 30 True 0.98 0.254588\n16 50 False 0.94 0.268659\n17 50 True 0.96 0.269792\n18 20 False 0.96 0.249941\n19 20 True 0.96 0.247138\n20 60 False 0.97 0.251301\n21 60 True 0.98 0.256504\n22 100 False 0.98 0.270049\n23 100 True 0.97 0.248972\n24 50 False 0.97 0.306356\n25 50 True 0.98 0.257544\n26 150 False 0.98 0.238811\n27 150 True 0.99 0.263939\n28 250 False 0.99 0.256558\n29 250 True 1.00 0.335823"
},
"execution_count": 22,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"df = pd.DataFrame(results)\n",
"df[[\"candidates\", \"rescore\", \"accuracy\", \"time\"]]\n",
"# df.to_csv(\"candidates-rescore-time.csv\", index=False)"
]
}
+24 -20
View File
@@ -11,23 +11,6 @@ from tqdm import tqdm
from loguru import logger
def locate_model_file(model_dir: Path, file_names: List[str]) -> Path:
"""
Find model path for both TransformerJS style `onnx` subdirectory structure and direct model weights structure used
by Optimum and Qdrant
"""
if not model_dir.is_dir():
raise ValueError(f"Provided model path '{model_dir}' is not a directory.")
for file_name in file_names:
file_paths = [path for path in model_dir.rglob(file_name) if path.is_file()]
if file_paths:
return file_paths[0]
raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
class ModelManagement:
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -104,20 +87,33 @@ class ModelManagement:
@classmethod
def download_files_from_huggingface(
cls, hf_source_repo: str, cache_dir: Optional[str] = None
cls,
hf_source_repo: str,
cache_dir: Optional[str] = None,
extra_patterns: Optional[List[str]] = None,
) -> str:
"""
Downloads a model from HuggingFace Hub.
Args:
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
cache_dir (Optional[str]): The path to the cache directory.
extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
includes the required model files.
Returns:
Path: The path to the model directory.
"""
allow_patterns = [
"config.json",
"tokenizer.json",
"tokenizer_config.json",
"special_tokens_map.json",
]
if extra_patterns is not None:
allow_patterns.extend(extra_patterns)
return snapshot_download(
repo_id=hf_source_repo,
ignore_patterns=["model.safetensors", "pytorch_model.bin"],
allow_patterns=allow_patterns,
cache_dir=cache_dir,
)
@@ -175,6 +171,9 @@ class ModelManagement:
model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
if model_tar_gz.exists():
model_tar_gz.unlink()
cls.download_file_from_gcs(
source_url,
output_path=str(model_tar_gz),
@@ -219,9 +218,14 @@ class ModelManagement:
url_source = model.get("sources", {}).get("url")
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))
cls.download_files_from_huggingface(
hf_source, cache_dir=str(cache_dir), extra_patterns=extra_patterns
)
)
except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
logger.error(
+9 -4
View File
@@ -6,11 +6,11 @@ from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, Ty
import numpy as np
import onnxruntime as ort
from fastembed.common.model_management import locate_model_file
from fastembed.common.models import load_tokenizer
from fastembed.common.utils import iter_batch
from fastembed.parallel_processor import ParallelWorkerPool, Worker
# Holds type of the embedding result
T = TypeVar("T")
@@ -34,8 +34,13 @@ class OnnxModel(Generic[T]):
"""
return onnx_input
def load_onnx_model(self, model_dir: Path, threads: Optional[int], max_length: int) -> None:
model_path = locate_model_file(model_dir, ["model.onnx", "model_optimized.onnx"])
def load_onnx_model(
self,
model_dir: Path,
model_file: str,
threads: Optional[int],
) -> None:
model_path = model_dir / model_file
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
onnx_providers = ["CPUExecutionProvider"]
@@ -47,7 +52,7 @@ class OnnxModel(Generic[T]):
so.intra_op_num_threads = threads
so.inter_op_num_threads = threads
self.tokenizer = load_tokenizer(model_dir=model_dir, max_length=max_length)
self.tokenizer = load_tokenizer(model_dir=model_dir)
self.model = ort.InferenceSession(
str(model_path), providers=onnx_providers, sess_options=so
)
+11 -7
View File
@@ -15,6 +15,7 @@ supported_splade_models = [
"sources": {
"hf": "Qdrant/SPLADE_PP_en_v1",
},
"model_file": "model.onnx",
},
{
"model": "prithivida/Splade_PP_en_v1",
@@ -24,6 +25,7 @@ supported_splade_models = [
"sources": {
"hf": "Qdrant/SPLADE_PP_en_v1",
},
"model_file": "model.onnx",
},
]
@@ -77,14 +79,16 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
super().__init__(model_name, cache_dir, threads, **kwargs)
self.model_name = model_name
self._model_description = self._get_model_description(model_name)
model_description = self._get_model_description(model_name)
cache_dir = define_cache_dir(cache_dir)
self._cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(self._model_description, self._cache_dir)
self._max_length = 512
model_dir = self.download_model(model_description, cache_dir)
self.load_onnx_model(self._model_dir, self.threads, self._max_length)
self.load_onnx_model(
model_dir=model_dir,
model_file=model_description["model_file"],
threads=threads,
)
def embed(
self,
@@ -110,7 +114,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
"""
yield from self._embed_documents(
model_name=self.model_name,
cache_dir=str(self._cache_dir),
cache_dir=str(self.cache_dir),
documents=documents,
batch_size=batch_size,
parallel=parallel,
+3
View File
@@ -15,6 +15,8 @@ supported_multilingual_e5_models = [
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
"hf": "qdrant/multilingual-e5-large-onnx",
},
"model_file": "model.onnx",
"additional_files": ["model.onnx_data"],
},
{
"model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
@@ -24,6 +26,7 @@ supported_multilingual_e5_models = [
"sources": {
"hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
},
"model_file": "onnx/model.onnx",
},
]
+2
View File
@@ -13,6 +13,7 @@ supported_jina_models = [
"description": "English embedding model supporting 8192 sequence length",
"size_in_GB": 0.52,
"sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
"model_file": "onnx/model.onnx",
},
{
"model": "jinaai/jina-embeddings-v2-small-en",
@@ -20,6 +21,7 @@ supported_jina_models = [
"description": "English embedding model supporting 8192 sequence length",
"size_in_GB": 0.12,
"sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
"model_file": "onnx/model.onnx",
},
]
+81 -32
View File
@@ -16,6 +16,7 @@ supported_onnx_models = [
"sources": {
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
},
"model_file": "model_optimized.onnx",
},
{
"model": "BAAI/bge-base-en-v1.5",
@@ -26,6 +27,7 @@ supported_onnx_models = [
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
"hf": "qdrant/bge-base-en-v1.5-onnx-q",
},
"model_file": "model_optimized.onnx",
},
{
"model": "BAAI/bge-large-en-v1.5",
@@ -35,6 +37,7 @@ supported_onnx_models = [
"sources": {
"hf": "qdrant/bge-large-en-v1.5-onnx",
},
"model_file": "model.onnx",
},
{
"model": "BAAI/bge-small-en",
@@ -44,18 +47,8 @@ supported_onnx_models = [
"sources": {
"url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
},
"model_file": "model_optimized.onnx",
},
# {
# "model": "BAAI/bge-small-en",
# "dim": 384,
# "description": "Fast English model",
# "size_in_GB": 0.2,
# "hf_sources": [],
# "compressed_url_sources": [
# "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
# "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
# ]
# },
{
"model": "BAAI/bge-small-en-v1.5",
"dim": 384,
@@ -64,6 +57,7 @@ supported_onnx_models = [
"sources": {
"hf": "qdrant/bge-small-en-v1.5-onnx-q",
},
"model_file": "model_optimized.onnx",
},
{
"model": "BAAI/bge-small-zh-v1.5",
@@ -73,6 +67,7 @@ supported_onnx_models = [
"sources": {
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
},
"model_file": "model_optimized.onnx",
},
{
"model": "sentence-transformers/all-MiniLM-L6-v2",
@@ -83,6 +78,7 @@ supported_onnx_models = [
"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": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
@@ -92,6 +88,7 @@ supported_onnx_models = [
"sources": {
"hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
},
"model_file": "model_optimized.onnx",
},
{
"model": "nomic-ai/nomic-embed-text-v1",
@@ -101,6 +98,7 @@ supported_onnx_models = [
"sources": {
"hf": "nomic-ai/nomic-embed-text-v1",
},
"model_file": "onnx/model.onnx",
},
{
"model": "nomic-ai/nomic-embed-text-v1.5",
@@ -110,6 +108,17 @@ supported_onnx_models = [
"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",
@@ -119,20 +128,8 @@ supported_onnx_models = [
"sources": {
"hf": "qdrant/gte-large-onnx",
},
"model_file": "model.onnx",
},
# {
# "model": "sentence-transformers/all-MiniLM-L6-v2",
# "dim": 384,
# "description": "Sentence Transformer model, MiniLM-L6-v2",
# "size_in_GB": 0.09,
# "hf_sources": [
# "qdrant/all-MiniLM-L6-v2-onnx"
# ],
# "compressed_url_sources": [
# "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
# "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
# ]
# }
{
"model": "mixedbread-ai/mxbai-embed-large-v1",
"dim": 1024,
@@ -141,6 +138,57 @@ supported_onnx_models = [
"sources": {
"hf": "mixedbread-ai/mxbai-embed-large-v1",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-xs",
"dim": 384,
"description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
"size_in_GB": 0.09,
"sources": {
"hf": "snowflake/snowflake-arctic-embed-xs",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-s",
"dim": 384,
"description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
"size_in_GB": 0.13,
"sources": {
"hf": "snowflake/snowflake-arctic-embed-s",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-m",
"dim": 768,
"description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
"size_in_GB": 0.43,
"sources": {
"hf": "Snowflake/snowflake-arctic-embed-m",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-m-long",
"dim": 768,
"description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
"size_in_GB": 0.54,
"sources": {
"hf": "snowflake/snowflake-arctic-embed-m-long",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-l",
"dim": 1024,
"description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
"size_in_GB": 1.02,
"sources": {
"hf": "snowflake/snowflake-arctic-embed-l",
},
"model_file": "onnx/model.onnx",
},
]
@@ -179,14 +227,15 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
super().__init__(model_name, cache_dir, threads, **kwargs)
self.model_name = model_name
self._model_description = self._get_model_description(model_name)
model_description = self._get_model_description(model_name)
cache_dir = define_cache_dir(cache_dir)
model_dir = self.download_model(model_description, cache_dir)
self._cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(self._model_description, self._cache_dir)
self._max_length = 512
self.load_onnx_model(self._model_dir, self.threads, self._max_length)
self.load_onnx_model(
model_dir=model_dir,
model_file=model_description["model_file"],
threads=threads,
)
def embed(
self,
@@ -212,7 +261,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
"""
yield from self._embed_documents(
model_name=self.model_name,
cache_dir=str(self._cache_dir),
cache_dir=str(self.cache_dir),
documents=documents,
batch_size=batch_size,
parallel=parallel,
Generated
+85 -73
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
[[package]]
name = "anyio"
@@ -876,13 +876,13 @@ license = ["ukkonen"]
[[package]]
name = "idna"
version = "3.6"
version = "3.7"
description = "Internationalized Domain Names in Applications (IDNA)"
optional = false
python-versions = ">=3.5"
files = [
{file = "idna-3.6-py3-none-any.whl", hash = "sha256:c05567e9c24a6b9faaa835c4821bad0590fbb9d5779e7caa6e1cc4978e7eb24f"},
{file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"},
{file = "idna-3.7-py3-none-any.whl", hash = "sha256:82fee1fc78add43492d3a1898bfa6d8a904cc97d8427f683ed8e798d07761aa0"},
{file = "idna-3.7.tar.gz", hash = "sha256:028ff3aadf0609c1fd278d8ea3089299412a7a8b9bd005dd08b9f8285bcb5cfc"},
]
[[package]]
@@ -2035,79 +2035,80 @@ files = [
[[package]]
name = "pillow"
version = "10.2.0"
version = "10.3.0"
description = "Python Imaging Library (Fork)"
optional = false
python-versions = ">=3.8"
files = [
{file = "pillow-10.2.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:7823bdd049099efa16e4246bdf15e5a13dbb18a51b68fa06d6c1d4d8b99a796e"},
{file = "pillow-10.2.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:83b2021f2ade7d1ed556bc50a399127d7fb245e725aa0113ebd05cfe88aaf588"},
{file = "pillow-10.2.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6fad5ff2f13d69b7e74ce5b4ecd12cc0ec530fcee76356cac6742785ff71c452"},
{file = "pillow-10.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:da2b52b37dad6d9ec64e653637a096905b258d2fc2b984c41ae7d08b938a67e4"},
{file = "pillow-10.2.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:47c0995fc4e7f79b5cfcab1fc437ff2890b770440f7696a3ba065ee0fd496563"},
{file = "pillow-10.2.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:322bdf3c9b556e9ffb18f93462e5f749d3444ce081290352c6070d014c93feb2"},
{file = "pillow-10.2.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:51f1a1bffc50e2e9492e87d8e09a17c5eea8409cda8d3f277eb6edc82813c17c"},
{file = "pillow-10.2.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:69ffdd6120a4737710a9eee73e1d2e37db89b620f702754b8f6e62594471dee0"},
{file = "pillow-10.2.0-cp310-cp310-win32.whl", hash = "sha256:c6dafac9e0f2b3c78df97e79af707cdc5ef8e88208d686a4847bab8266870023"},
{file = "pillow-10.2.0-cp310-cp310-win_amd64.whl", hash = "sha256:aebb6044806f2e16ecc07b2a2637ee1ef67a11840a66752751714a0d924adf72"},
{file = "pillow-10.2.0-cp310-cp310-win_arm64.whl", hash = "sha256:7049e301399273a0136ff39b84c3678e314f2158f50f517bc50285fb5ec847ad"},
{file = "pillow-10.2.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:35bb52c37f256f662abdfa49d2dfa6ce5d93281d323a9af377a120e89a9eafb5"},
{file = "pillow-10.2.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9c23f307202661071d94b5e384e1e1dc7dfb972a28a2310e4ee16103e66ddb67"},
{file = "pillow-10.2.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:773efe0603db30c281521a7c0214cad7836c03b8ccff897beae9b47c0b657d61"},
{file = "pillow-10.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:11fa2e5984b949b0dd6d7a94d967743d87c577ff0b83392f17cb3990d0d2fd6e"},
{file = "pillow-10.2.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:716d30ed977be8b37d3ef185fecb9e5a1d62d110dfbdcd1e2a122ab46fddb03f"},
{file = "pillow-10.2.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a086c2af425c5f62a65e12fbf385f7c9fcb8f107d0849dba5839461a129cf311"},
{file = "pillow-10.2.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:c8de2789052ed501dd829e9cae8d3dcce7acb4777ea4a479c14521c942d395b1"},
{file = "pillow-10.2.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:609448742444d9290fd687940ac0b57fb35e6fd92bdb65386e08e99af60bf757"},
{file = "pillow-10.2.0-cp311-cp311-win32.whl", hash = "sha256:823ef7a27cf86df6597fa0671066c1b596f69eba53efa3d1e1cb8b30f3533068"},
{file = "pillow-10.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:1da3b2703afd040cf65ec97efea81cfba59cdbed9c11d8efc5ab09df9509fc56"},
{file = "pillow-10.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:edca80cbfb2b68d7b56930b84a0e45ae1694aeba0541f798e908a49d66b837f1"},
{file = "pillow-10.2.0-cp312-cp312-macosx_10_10_x86_64.whl", hash = "sha256:1b5e1b74d1bd1b78bc3477528919414874748dd363e6272efd5abf7654e68bef"},
{file = "pillow-10.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0eae2073305f451d8ecacb5474997c08569fb4eb4ac231ffa4ad7d342fdc25ac"},
{file = "pillow-10.2.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b7c2286c23cd350b80d2fc9d424fc797575fb16f854b831d16fd47ceec078f2c"},
{file = "pillow-10.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1e23412b5c41e58cec602f1135c57dfcf15482013ce6e5f093a86db69646a5aa"},
{file = "pillow-10.2.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:52a50aa3fb3acb9cf7213573ef55d31d6eca37f5709c69e6858fe3bc04a5c2a2"},
{file = "pillow-10.2.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:127cee571038f252a552760076407f9cff79761c3d436a12af6000cd182a9d04"},
{file = "pillow-10.2.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:8d12251f02d69d8310b046e82572ed486685c38f02176bd08baf216746eb947f"},
{file = "pillow-10.2.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:54f1852cd531aa981bc0965b7d609f5f6cc8ce8c41b1139f6ed6b3c54ab82bfb"},
{file = "pillow-10.2.0-cp312-cp312-win32.whl", hash = "sha256:257d8788df5ca62c980314053197f4d46eefedf4e6175bc9412f14412ec4ea2f"},
{file = "pillow-10.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:154e939c5f0053a383de4fd3d3da48d9427a7e985f58af8e94d0b3c9fcfcf4f9"},
{file = "pillow-10.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:f379abd2f1e3dddb2b61bc67977a6b5a0a3f7485538bcc6f39ec76163891ee48"},
{file = "pillow-10.2.0-cp38-cp38-macosx_10_10_x86_64.whl", hash = "sha256:8373c6c251f7ef8bda6675dd6d2b3a0fcc31edf1201266b5cf608b62a37407f9"},
{file = "pillow-10.2.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:870ea1ada0899fd0b79643990809323b389d4d1d46c192f97342eeb6ee0b8483"},
{file = "pillow-10.2.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b4b6b1e20608493548b1f32bce8cca185bf0480983890403d3b8753e44077129"},
{file = "pillow-10.2.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3031709084b6e7852d00479fd1d310b07d0ba82765f973b543c8af5061cf990e"},
{file = "pillow-10.2.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:3ff074fc97dd4e80543a3e91f69d58889baf2002b6be64347ea8cf5533188213"},
{file = "pillow-10.2.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:cb4c38abeef13c61d6916f264d4845fab99d7b711be96c326b84df9e3e0ff62d"},
{file = "pillow-10.2.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:b1b3020d90c2d8e1dae29cf3ce54f8094f7938460fb5ce8bc5c01450b01fbaf6"},
{file = "pillow-10.2.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:170aeb00224ab3dc54230c797f8404507240dd868cf52066f66a41b33169bdbe"},
{file = "pillow-10.2.0-cp38-cp38-win32.whl", hash = "sha256:c4225f5220f46b2fde568c74fca27ae9771536c2e29d7c04f4fb62c83275ac4e"},
{file = "pillow-10.2.0-cp38-cp38-win_amd64.whl", hash = "sha256:0689b5a8c5288bc0504d9fcee48f61a6a586b9b98514d7d29b840143d6734f39"},
{file = "pillow-10.2.0-cp39-cp39-macosx_10_10_x86_64.whl", hash = "sha256:b792a349405fbc0163190fde0dc7b3fef3c9268292586cf5645598b48e63dc67"},
{file = "pillow-10.2.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c570f24be1e468e3f0ce7ef56a89a60f0e05b30a3669a459e419c6eac2c35364"},
{file = "pillow-10.2.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d8ecd059fdaf60c1963c58ceb8997b32e9dc1b911f5da5307aab614f1ce5c2fb"},
{file = "pillow-10.2.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c365fd1703040de1ec284b176d6af5abe21b427cb3a5ff68e0759e1e313a5e7e"},
{file = "pillow-10.2.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:70c61d4c475835a19b3a5aa42492409878bbca7438554a1f89d20d58a7c75c01"},
{file = "pillow-10.2.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:b6f491cdf80ae540738859d9766783e3b3c8e5bd37f5dfa0b76abdecc5081f13"},
{file = "pillow-10.2.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:9d189550615b4948f45252d7f005e53c2040cea1af5b60d6f79491a6e147eef7"},
{file = "pillow-10.2.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:49d9ba1ed0ef3e061088cd1e7538a0759aab559e2e0a80a36f9fd9d8c0c21591"},
{file = "pillow-10.2.0-cp39-cp39-win32.whl", hash = "sha256:babf5acfede515f176833ed6028754cbcd0d206f7f614ea3447d67c33be12516"},
{file = "pillow-10.2.0-cp39-cp39-win_amd64.whl", hash = "sha256:0304004f8067386b477d20a518b50f3fa658a28d44e4116970abfcd94fac34a8"},
{file = "pillow-10.2.0-cp39-cp39-win_arm64.whl", hash = "sha256:0fb3e7fc88a14eacd303e90481ad983fd5b69c761e9e6ef94c983f91025da869"},
{file = "pillow-10.2.0-pp310-pypy310_pp73-macosx_10_10_x86_64.whl", hash = "sha256:322209c642aabdd6207517e9739c704dc9f9db943015535783239022002f054a"},
{file = "pillow-10.2.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3eedd52442c0a5ff4f887fab0c1c0bb164d8635b32c894bc1faf4c618dd89df2"},
{file = "pillow-10.2.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cb28c753fd5eb3dd859b4ee95de66cc62af91bcff5db5f2571d32a520baf1f04"},
{file = "pillow-10.2.0-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:33870dc4653c5017bf4c8873e5488d8f8d5f8935e2f1fb9a2208c47cdd66efd2"},
{file = "pillow-10.2.0-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:3c31822339516fb3c82d03f30e22b1d038da87ef27b6a78c9549888f8ceda39a"},
{file = "pillow-10.2.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:a2b56ba36e05f973d450582fb015594aaa78834fefe8dfb8fcd79b93e64ba4c6"},
{file = "pillow-10.2.0-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:d8e6aeb9201e655354b3ad049cb77d19813ad4ece0df1249d3c793de3774f8c7"},
{file = "pillow-10.2.0-pp39-pypy39_pp73-macosx_10_10_x86_64.whl", hash = "sha256:2247178effb34a77c11c0e8ac355c7a741ceca0a732b27bf11e747bbc950722f"},
{file = "pillow-10.2.0-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:15587643b9e5eb26c48e49a7b33659790d28f190fc514a322d55da2fb5c2950e"},
{file = "pillow-10.2.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:753cd8f2086b2b80180d9b3010dd4ed147efc167c90d3bf593fe2af21265e5a5"},
{file = "pillow-10.2.0-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:7c8f97e8e7a9009bcacbe3766a36175056c12f9a44e6e6f2d5caad06dcfbf03b"},
{file = "pillow-10.2.0-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:d1b35bcd6c5543b9cb547dee3150c93008f8dd0f1fef78fc0cd2b141c5baf58a"},
{file = "pillow-10.2.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:fe4c15f6c9285dc54ce6553a3ce908ed37c8f3825b5a51a15c91442bb955b868"},
{file = "pillow-10.2.0.tar.gz", hash = "sha256:e87f0b2c78157e12d7686b27d63c070fd65d994e8ddae6f328e0dcf4a0cd007e"},
{file = "pillow-10.3.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:90b9e29824800e90c84e4022dd5cc16eb2d9605ee13f05d47641eb183cd73d45"},
{file = "pillow-10.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a2c405445c79c3f5a124573a051062300936b0281fee57637e706453e452746c"},
{file = "pillow-10.3.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:78618cdbccaa74d3f88d0ad6cb8ac3007f1a6fa5c6f19af64b55ca170bfa1edf"},
{file = "pillow-10.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:261ddb7ca91fcf71757979534fb4c128448b5b4c55cb6152d280312062f69599"},
{file = "pillow-10.3.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:ce49c67f4ea0609933d01c0731b34b8695a7a748d6c8d186f95e7d085d2fe475"},
{file = "pillow-10.3.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:b14f16f94cbc61215115b9b1236f9c18403c15dd3c52cf629072afa9d54c1cbf"},
{file = "pillow-10.3.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:d33891be6df59d93df4d846640f0e46f1a807339f09e79a8040bc887bdcd7ed3"},
{file = "pillow-10.3.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:b50811d664d392f02f7761621303eba9d1b056fb1868c8cdf4231279645c25f5"},
{file = "pillow-10.3.0-cp310-cp310-win32.whl", hash = "sha256:ca2870d5d10d8726a27396d3ca4cf7976cec0f3cb706debe88e3a5bd4610f7d2"},
{file = "pillow-10.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:f0d0591a0aeaefdaf9a5e545e7485f89910c977087e7de2b6c388aec32011e9f"},
{file = "pillow-10.3.0-cp310-cp310-win_arm64.whl", hash = "sha256:ccce24b7ad89adb5a1e34a6ba96ac2530046763912806ad4c247356a8f33a67b"},
{file = "pillow-10.3.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:5f77cf66e96ae734717d341c145c5949c63180842a545c47a0ce7ae52ca83795"},
{file = "pillow-10.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e4b878386c4bf293578b48fc570b84ecfe477d3b77ba39a6e87150af77f40c57"},
{file = "pillow-10.3.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fdcbb4068117dfd9ce0138d068ac512843c52295ed996ae6dd1faf537b6dbc27"},
{file = "pillow-10.3.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9797a6c8fe16f25749b371c02e2ade0efb51155e767a971c61734b1bf6293994"},
{file = "pillow-10.3.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:9e91179a242bbc99be65e139e30690e081fe6cb91a8e77faf4c409653de39451"},
{file = "pillow-10.3.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:1b87bd9d81d179bd8ab871603bd80d8645729939f90b71e62914e816a76fc6bd"},
{file = "pillow-10.3.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:81d09caa7b27ef4e61cb7d8fbf1714f5aec1c6b6c5270ee53504981e6e9121ad"},
{file = "pillow-10.3.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:048ad577748b9fa4a99a0548c64f2cb8d672d5bf2e643a739ac8faff1164238c"},
{file = "pillow-10.3.0-cp311-cp311-win32.whl", hash = "sha256:7161ec49ef0800947dc5570f86568a7bb36fa97dd09e9827dc02b718c5643f09"},
{file = "pillow-10.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:8eb0908e954d093b02a543dc963984d6e99ad2b5e36503d8a0aaf040505f747d"},
{file = "pillow-10.3.0-cp311-cp311-win_arm64.whl", hash = "sha256:4e6f7d1c414191c1199f8996d3f2282b9ebea0945693fb67392c75a3a320941f"},
{file = "pillow-10.3.0-cp312-cp312-macosx_10_10_x86_64.whl", hash = "sha256:e46f38133e5a060d46bd630faa4d9fa0202377495df1f068a8299fd78c84de84"},
{file = "pillow-10.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:50b8eae8f7334ec826d6eeffaeeb00e36b5e24aa0b9df322c247539714c6df19"},
{file = "pillow-10.3.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d3bea1c75f8c53ee4d505c3e67d8c158ad4df0d83170605b50b64025917f338"},
{file = "pillow-10.3.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:19aeb96d43902f0a783946a0a87dbdad5c84c936025b8419da0a0cd7724356b1"},
{file = "pillow-10.3.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:74d28c17412d9caa1066f7a31df8403ec23d5268ba46cd0ad2c50fb82ae40462"},
{file = "pillow-10.3.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ff61bfd9253c3915e6d41c651d5f962da23eda633cf02262990094a18a55371a"},
{file = "pillow-10.3.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d886f5d353333b4771d21267c7ecc75b710f1a73d72d03ca06df49b09015a9ef"},
{file = "pillow-10.3.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4b5ec25d8b17217d635f8935dbc1b9aa5907962fae29dff220f2659487891cd3"},
{file = "pillow-10.3.0-cp312-cp312-win32.whl", hash = "sha256:51243f1ed5161b9945011a7360e997729776f6e5d7005ba0c6879267d4c5139d"},
{file = "pillow-10.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:412444afb8c4c7a6cc11a47dade32982439925537e483be7c0ae0cf96c4f6a0b"},
{file = "pillow-10.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:798232c92e7665fe82ac085f9d8e8ca98826f8e27859d9a96b41d519ecd2e49a"},
{file = "pillow-10.3.0-cp38-cp38-macosx_10_10_x86_64.whl", hash = "sha256:4eaa22f0d22b1a7e93ff0a596d57fdede2e550aecffb5a1ef1106aaece48e96b"},
{file = "pillow-10.3.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:cd5e14fbf22a87321b24c88669aad3a51ec052eb145315b3da3b7e3cc105b9a2"},
{file = "pillow-10.3.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1530e8f3a4b965eb6a7785cf17a426c779333eb62c9a7d1bbcf3ffd5bf77a4aa"},
{file = "pillow-10.3.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d512aafa1d32efa014fa041d38868fda85028e3f930a96f85d49c7d8ddc0383"},
{file = "pillow-10.3.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:339894035d0ede518b16073bdc2feef4c991ee991a29774b33e515f1d308e08d"},
{file = "pillow-10.3.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:aa7e402ce11f0885305bfb6afb3434b3cd8f53b563ac065452d9d5654c7b86fd"},
{file = "pillow-10.3.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:0ea2a783a2bdf2a561808fe4a7a12e9aa3799b701ba305de596bc48b8bdfce9d"},
{file = "pillow-10.3.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:c78e1b00a87ce43bb37642c0812315b411e856a905d58d597750eb79802aaaa3"},
{file = "pillow-10.3.0-cp38-cp38-win32.whl", hash = "sha256:72d622d262e463dfb7595202d229f5f3ab4b852289a1cd09650362db23b9eb0b"},
{file = "pillow-10.3.0-cp38-cp38-win_amd64.whl", hash = "sha256:2034f6759a722da3a3dbd91a81148cf884e91d1b747992ca288ab88c1de15999"},
{file = "pillow-10.3.0-cp39-cp39-macosx_10_10_x86_64.whl", hash = "sha256:2ed854e716a89b1afcedea551cd85f2eb2a807613752ab997b9974aaa0d56936"},
{file = "pillow-10.3.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:dc1a390a82755a8c26c9964d457d4c9cbec5405896cba94cf51f36ea0d855002"},
{file = "pillow-10.3.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4203efca580f0dd6f882ca211f923168548f7ba334c189e9eab1178ab840bf60"},
{file = "pillow-10.3.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3102045a10945173d38336f6e71a8dc71bcaeed55c3123ad4af82c52807b9375"},
{file = "pillow-10.3.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:6fb1b30043271ec92dc65f6d9f0b7a830c210b8a96423074b15c7bc999975f57"},
{file = "pillow-10.3.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:1dfc94946bc60ea375cc39cff0b8da6c7e5f8fcdc1d946beb8da5c216156ddd8"},
{file = "pillow-10.3.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:b09b86b27a064c9624d0a6c54da01c1beaf5b6cadfa609cf63789b1d08a797b9"},
{file = "pillow-10.3.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:d3b2348a78bc939b4fed6552abfd2e7988e0f81443ef3911a4b8498ca084f6eb"},
{file = "pillow-10.3.0-cp39-cp39-win32.whl", hash = "sha256:45ebc7b45406febf07fef35d856f0293a92e7417ae7933207e90bf9090b70572"},
{file = "pillow-10.3.0-cp39-cp39-win_amd64.whl", hash = "sha256:0ba26351b137ca4e0db0342d5d00d2e355eb29372c05afd544ebf47c0956ffeb"},
{file = "pillow-10.3.0-cp39-cp39-win_arm64.whl", hash = "sha256:50fd3f6b26e3441ae07b7c979309638b72abc1a25da31a81a7fbd9495713ef4f"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-macosx_10_10_x86_64.whl", hash = "sha256:6b02471b72526ab8a18c39cb7967b72d194ec53c1fd0a70b050565a0f366d355"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:8ab74c06ffdab957d7670c2a5a6e1a70181cd10b727cd788c4dd9005b6a8acd9"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:048eeade4c33fdf7e08da40ef402e748df113fd0b4584e32c4af74fe78baaeb2"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e2ec1e921fd07c7cda7962bad283acc2f2a9ccc1b971ee4b216b75fad6f0463"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:4c8e73e99da7db1b4cad7f8d682cf6abad7844da39834c288fbfa394a47bbced"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:16563993329b79513f59142a6b02055e10514c1a8e86dca8b48a893e33cf91e3"},
{file = "pillow-10.3.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:dd78700f5788ae180b5ee8902c6aea5a5726bac7c364b202b4b3e3ba2d293170"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-macosx_10_10_x86_64.whl", hash = "sha256:aff76a55a8aa8364d25400a210a65ff59d0168e0b4285ba6bf2bd83cf675ba32"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-macosx_11_0_arm64.whl", hash = "sha256:b7bc2176354defba3edc2b9a777744462da2f8e921fbaf61e52acb95bafa9828"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:793b4e24db2e8742ca6423d3fde8396db336698c55cd34b660663ee9e45ed37f"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d93480005693d247f8346bc8ee28c72a2191bdf1f6b5db469c096c0c867ac015"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:c83341b89884e2b2e55886e8fbbf37c3fa5efd6c8907124aeb72f285ae5696e5"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:1a1d1915db1a4fdb2754b9de292642a39a7fb28f1736699527bb649484fb966a"},
{file = "pillow-10.3.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:a0eaa93d054751ee9964afa21c06247779b90440ca41d184aeb5d410f20ff591"},
{file = "pillow-10.3.0.tar.gz", hash = "sha256:9d2455fbf44c914840c793e89aa82d0e1763a14253a000743719ae5946814b2d"},
]
[package.extras]
@@ -2441,6 +2442,7 @@ files = [
{file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"},
{file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"},
{file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"},
{file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"},
{file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"},
{file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"},
{file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"},
@@ -2448,8 +2450,16 @@ files = [
{file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"},
{file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"},
{file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"},
{file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"},
{file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"},
{file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
{file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
{file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
{file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"},
{file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"},
{file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"},
{file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"},
@@ -2466,6 +2476,7 @@ files = [
{file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"},
{file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"},
{file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"},
{file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"},
{file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"},
{file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"},
{file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"},
@@ -2473,6 +2484,7 @@ files = [
{file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"},
{file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"},
{file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"},
{file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"},
{file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"},
{file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"},
{file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
+12 -2
View File
@@ -26,16 +26,26 @@ CANONICAL_VECTOR_VALUES = {
"nomic-ai/nomic-embed-text-v1.5": np.array(
[-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
),
"nomic-ai/nomic-embed-text-v1.5-Q": np.array(
[-0.01554983, 0.0129992 , -0.17909265, -0.01062993, 0.00512859]
),
"thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
"mixedbread-ai/mxbai-embed-large-v1": np.array([0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]),
"snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
"snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
"snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
"snowflake/snowflake-arctic-embed-m-long": np.array(
[0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
),
"snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
}
def test_embedding():
is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
is_ci = os.getenv("CI")
for model_desc in TextEmbedding.list_supported_models():
if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
if not is_ci and model_desc["size_in_GB"] > 1:
continue
dim = model_desc["dim"]