fastembed-gpu 0.2.7__tar.gz → 0.3.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/PKG-INFO +114 -15
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/README.md +103 -8
- fastembed_gpu-0.3.1/fastembed/__init__.py +20 -0
- fastembed_gpu-0.3.1/fastembed/common/__init__.py +3 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/common/model_management.py +22 -8
- fastembed_gpu-0.3.1/fastembed/common/onnx_model.py +123 -0
- fastembed_gpu-0.2.7/fastembed/common/models.py → fastembed_gpu-0.3.1/fastembed/common/preprocessor_utils.py +39 -17
- fastembed_gpu-0.3.1/fastembed/common/types.py +14 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/common/utils.py +11 -1
- fastembed_gpu-0.3.1/fastembed/image/__init__.py +3 -0
- fastembed_gpu-0.3.1/fastembed/image/image_embedding.py +94 -0
- fastembed_gpu-0.3.1/fastembed/image/image_embedding_base.py +44 -0
- fastembed_gpu-0.3.1/fastembed/image/onnx_embedding.py +131 -0
- fastembed_gpu-0.3.1/fastembed/image/onnx_image_model.py +108 -0
- fastembed_gpu-0.3.1/fastembed/image/transform/functional.py +124 -0
- fastembed_gpu-0.3.1/fastembed/image/transform/operators.py +198 -0
- fastembed_gpu-0.3.1/fastembed/late_interaction/__init__.py +5 -0
- fastembed_gpu-0.3.1/fastembed/late_interaction/colbert.py +194 -0
- fastembed_gpu-0.3.1/fastembed/late_interaction/late_interaction_embedding_base.py +62 -0
- fastembed_gpu-0.3.1/fastembed/late_interaction/late_interaction_text_embedding.py +109 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/parallel_processor.py +6 -2
- fastembed_gpu-0.3.1/fastembed/sparse/bm25.py +284 -0
- fastembed_gpu-0.3.1/fastembed/sparse/bm42.py +292 -0
- fastembed_gpu-0.3.1/fastembed/sparse/sparse_embedding_base.py +85 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/sparse/sparse_text_embedding.py +31 -7
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/sparse/splade_pp.py +18 -21
- fastembed_gpu-0.3.1/fastembed/sparse/utils/tokenizer.py +113 -0
- fastembed_gpu-0.3.1/fastembed/text/clip_embedding.py +49 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/e5_onnx_embedding.py +10 -8
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/jina_onnx_embedding.py +21 -12
- fastembed_gpu-0.3.1/fastembed/text/mini_lm_embedding.py +58 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/onnx_embedding.py +18 -27
- fastembed_gpu-0.3.1/fastembed/text/onnx_text_model.py +126 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/text_embedding.py +14 -3
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/text_embedding_base.py +3 -1
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/pyproject.toml +11 -7
- fastembed_gpu-0.2.7/fastembed/__init__.py +0 -12
- fastembed_gpu-0.2.7/fastembed/common/__init__.py +0 -3
- fastembed_gpu-0.2.7/fastembed/common/onnx_model.py +0 -165
- fastembed_gpu-0.2.7/fastembed/image/__init__.py +0 -0
- fastembed_gpu-0.2.7/fastembed/sparse/sparse_embedding_base.py +0 -44
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/LICENSE +0 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/embedding.py +0 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/sparse/__init__.py +0 -0
- {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/__init__.py +0 -0
|
@@ -1,12 +1,12 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: fastembed-gpu
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.1
|
|
4
4
|
Summary: Fast, light, accurate library built for retrieval embedding generation
|
|
5
5
|
Home-page: https://github.com/qdrant/fastembed
|
|
6
6
|
License: Apache License
|
|
7
7
|
Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
|
|
8
|
-
Author:
|
|
9
|
-
Author-email:
|
|
8
|
+
Author: Qdrant Team
|
|
9
|
+
Author-email: info@qdrant.tech
|
|
10
10
|
Requires-Python: >=3.8.0,<3.13
|
|
11
11
|
Classifier: License :: Other/Proprietary License
|
|
12
12
|
Classifier: Programming Language :: Python :: 3
|
|
@@ -15,13 +15,17 @@ Classifier: Programming Language :: Python :: 3.9
|
|
|
15
15
|
Classifier: Programming Language :: Python :: 3.10
|
|
16
16
|
Classifier: Programming Language :: Python :: 3.11
|
|
17
17
|
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
-
Requires-Dist:
|
|
18
|
+
Requires-Dist: PyStemmer (>=2.2.0,<3.0.0)
|
|
19
|
+
Requires-Dist: huggingface-hub (>=0.20,<1.0)
|
|
19
20
|
Requires-Dist: loguru (>=0.7.2,<0.8.0)
|
|
20
|
-
Requires-Dist:
|
|
21
|
-
Requires-Dist: numpy (>=1.
|
|
21
|
+
Requires-Dist: mmh3 (>=4.0,<5.0)
|
|
22
|
+
Requires-Dist: numpy (>=1.21,<2) ; python_version < "3.12"
|
|
23
|
+
Requires-Dist: numpy (>=1.26,<2) ; python_version >= "3.12"
|
|
22
24
|
Requires-Dist: onnxruntime-gpu (>=1.17.0,<2.0.0)
|
|
25
|
+
Requires-Dist: pillow (>=10.3.0,<11.0.0)
|
|
23
26
|
Requires-Dist: requests (>=2.31,<3.0)
|
|
24
|
-
Requires-Dist:
|
|
27
|
+
Requires-Dist: snowballstemmer (>=2.2.0,<3.0.0)
|
|
28
|
+
Requires-Dist: tokenizers (>=0.15,<1.0)
|
|
25
29
|
Requires-Dist: tqdm (>=4.66,<5.0)
|
|
26
30
|
Project-URL: Repository, https://github.com/qdrant/fastembed
|
|
27
31
|
Description-Content-Type: text/markdown
|
|
@@ -46,11 +50,9 @@ To install the FastEmbed library, pip works best. You can install it with or wit
|
|
|
46
50
|
|
|
47
51
|
```bash
|
|
48
52
|
pip install fastembed
|
|
49
|
-
```
|
|
50
53
|
|
|
51
|
-
|
|
54
|
+
# or with GPU support
|
|
52
55
|
|
|
53
|
-
```bash
|
|
54
56
|
pip install fastembed-gpu
|
|
55
57
|
```
|
|
56
58
|
|
|
@@ -76,19 +78,116 @@ embeddings_list = list(embedding_model.embed(documents))
|
|
|
76
78
|
len(embeddings_list[0]) # Vector of 384 dimensions
|
|
77
79
|
```
|
|
78
80
|
|
|
79
|
-
|
|
81
|
+
Fastembed supports a variety of models for different tasks and modalities.
|
|
82
|
+
The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
|
|
83
|
+
### 🎒 Dense text embeddings
|
|
84
|
+
|
|
85
|
+
```python
|
|
86
|
+
from fastembed import TextEmbedding
|
|
87
|
+
|
|
88
|
+
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
|
|
89
|
+
embeddings = list(embedding_model.embed(documents))
|
|
90
|
+
|
|
91
|
+
# [
|
|
92
|
+
# array([-0.1115, 0.0097, 0.0052, 0.0195, ...], dtype=float32),
|
|
93
|
+
# array([-0.1019, 0.0635, -0.0332, 0.0522, ...], dtype=float32)
|
|
94
|
+
# ]
|
|
95
|
+
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
### 🔱 Sparse text embeddings
|
|
101
|
+
|
|
102
|
+
* SPLADE++
|
|
103
|
+
|
|
104
|
+
```python
|
|
105
|
+
from fastembed import SparseTextEmbedding
|
|
106
|
+
|
|
107
|
+
model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
|
|
108
|
+
embeddings = list(embedding_model.embed(documents))
|
|
109
|
+
|
|
110
|
+
# [
|
|
111
|
+
# SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
|
|
112
|
+
# SparseEmbedding(indices=[ 38, 12, 91, ... ], values=[0.11, 0.22, 0.39, ...])
|
|
113
|
+
# ]
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
<!--
|
|
117
|
+
* BM42 - ([link](ToDo))
|
|
118
|
+
|
|
119
|
+
```
|
|
120
|
+
from fastembed import SparseTextEmbedding
|
|
121
|
+
|
|
122
|
+
model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
|
|
123
|
+
embeddings = list(embedding_model.embed(documents))
|
|
124
|
+
|
|
125
|
+
# [
|
|
126
|
+
# SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
|
|
127
|
+
# SparseEmbedding(indices=[ 38, 12, 91, ... ], values=[0.11, 0.22, 0.39, ...])
|
|
128
|
+
# ]
|
|
129
|
+
```
|
|
130
|
+
-->
|
|
131
|
+
|
|
132
|
+
### 🦥 Late interaction models (aka ColBERT)
|
|
80
133
|
|
|
81
|
-
|
|
82
|
-
|
|
134
|
+
|
|
135
|
+
```python
|
|
136
|
+
from fastembed import LateInteractionTextEmbedding
|
|
137
|
+
|
|
138
|
+
model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
|
|
139
|
+
embeddings = list(embedding_model.embed(documents))
|
|
140
|
+
|
|
141
|
+
# [
|
|
142
|
+
# array([
|
|
143
|
+
# [-0.1115, 0.0097, 0.0052, 0.0195, ...],
|
|
144
|
+
# [-0.1019, 0.0635, -0.0332, 0.0522, ...],
|
|
145
|
+
# ]),
|
|
146
|
+
# array([
|
|
147
|
+
# [-0.9019, 0.0335, -0.0032, 0.0991, ...],
|
|
148
|
+
# [-0.2115, 0.8097, 0.1052, 0.0195, ...],
|
|
149
|
+
# ]),
|
|
150
|
+
# ]
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
### 🖼️ Image embeddings
|
|
154
|
+
|
|
155
|
+
```python
|
|
156
|
+
from fastembed import ImageEmbedding
|
|
157
|
+
|
|
158
|
+
images = [
|
|
159
|
+
"./path/to/image1.jpg",
|
|
160
|
+
"./path/to/image2.jpg",
|
|
161
|
+
]
|
|
162
|
+
|
|
163
|
+
model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
|
|
164
|
+
embeddings = list(embedding_model.embed(images))
|
|
165
|
+
|
|
166
|
+
# [
|
|
167
|
+
# array([-0.1115, 0.0097, 0.0052, 0.0195, ...], dtype=float32),
|
|
168
|
+
# array([-0.1019, 0.0635, -0.0332, 0.0522, ...], dtype=float32)
|
|
169
|
+
# ]
|
|
170
|
+
```
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
## ⚡️ FastEmbed on a GPU
|
|
174
|
+
|
|
175
|
+
FastEmbed supports running on GPU devices.
|
|
176
|
+
It requires installation of the `fastembed-gpu` package.
|
|
83
177
|
|
|
84
178
|
```bash
|
|
85
179
|
pip install fastembed-gpu
|
|
86
|
-
```
|
|
180
|
+
```
|
|
181
|
+
|
|
182
|
+
Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
|
|
87
183
|
|
|
88
184
|
```python
|
|
89
185
|
from fastembed import TextEmbedding
|
|
90
186
|
|
|
91
|
-
embedding_model = TextEmbedding(
|
|
187
|
+
embedding_model = TextEmbedding(
|
|
188
|
+
model_name="BAAI/bge-small-en-v1.5",
|
|
189
|
+
providers=["CUDAExecutionProvider"]
|
|
190
|
+
)
|
|
92
191
|
print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
|
|
93
192
|
|
|
94
193
|
```
|
|
@@ -18,11 +18,9 @@ To install the FastEmbed library, pip works best. You can install it with or wit
|
|
|
18
18
|
|
|
19
19
|
```bash
|
|
20
20
|
pip install fastembed
|
|
21
|
-
```
|
|
22
21
|
|
|
23
|
-
|
|
22
|
+
# or with GPU support
|
|
24
23
|
|
|
25
|
-
```bash
|
|
26
24
|
pip install fastembed-gpu
|
|
27
25
|
```
|
|
28
26
|
|
|
@@ -48,19 +46,116 @@ embeddings_list = list(embedding_model.embed(documents))
|
|
|
48
46
|
len(embeddings_list[0]) # Vector of 384 dimensions
|
|
49
47
|
```
|
|
50
48
|
|
|
51
|
-
|
|
49
|
+
Fastembed supports a variety of models for different tasks and modalities.
|
|
50
|
+
The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
|
|
51
|
+
### 🎒 Dense text embeddings
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from fastembed import TextEmbedding
|
|
55
|
+
|
|
56
|
+
model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
|
|
57
|
+
embeddings = list(embedding_model.embed(documents))
|
|
58
|
+
|
|
59
|
+
# [
|
|
60
|
+
# array([-0.1115, 0.0097, 0.0052, 0.0195, ...], dtype=float32),
|
|
61
|
+
# array([-0.1019, 0.0635, -0.0332, 0.0522, ...], dtype=float32)
|
|
62
|
+
# ]
|
|
63
|
+
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
### 🔱 Sparse text embeddings
|
|
69
|
+
|
|
70
|
+
* SPLADE++
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
from fastembed import SparseTextEmbedding
|
|
74
|
+
|
|
75
|
+
model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
|
|
76
|
+
embeddings = list(embedding_model.embed(documents))
|
|
77
|
+
|
|
78
|
+
# [
|
|
79
|
+
# SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
|
|
80
|
+
# SparseEmbedding(indices=[ 38, 12, 91, ... ], values=[0.11, 0.22, 0.39, ...])
|
|
81
|
+
# ]
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
<!--
|
|
85
|
+
* BM42 - ([link](ToDo))
|
|
86
|
+
|
|
87
|
+
```
|
|
88
|
+
from fastembed import SparseTextEmbedding
|
|
89
|
+
|
|
90
|
+
model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
|
|
91
|
+
embeddings = list(embedding_model.embed(documents))
|
|
92
|
+
|
|
93
|
+
# [
|
|
94
|
+
# SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
|
|
95
|
+
# SparseEmbedding(indices=[ 38, 12, 91, ... ], values=[0.11, 0.22, 0.39, ...])
|
|
96
|
+
# ]
|
|
97
|
+
```
|
|
98
|
+
-->
|
|
99
|
+
|
|
100
|
+
### 🦥 Late interaction models (aka ColBERT)
|
|
52
101
|
|
|
53
|
-
|
|
54
|
-
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
from fastembed import LateInteractionTextEmbedding
|
|
105
|
+
|
|
106
|
+
model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
|
|
107
|
+
embeddings = list(embedding_model.embed(documents))
|
|
108
|
+
|
|
109
|
+
# [
|
|
110
|
+
# array([
|
|
111
|
+
# [-0.1115, 0.0097, 0.0052, 0.0195, ...],
|
|
112
|
+
# [-0.1019, 0.0635, -0.0332, 0.0522, ...],
|
|
113
|
+
# ]),
|
|
114
|
+
# array([
|
|
115
|
+
# [-0.9019, 0.0335, -0.0032, 0.0991, ...],
|
|
116
|
+
# [-0.2115, 0.8097, 0.1052, 0.0195, ...],
|
|
117
|
+
# ]),
|
|
118
|
+
# ]
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
### 🖼️ Image embeddings
|
|
122
|
+
|
|
123
|
+
```python
|
|
124
|
+
from fastembed import ImageEmbedding
|
|
125
|
+
|
|
126
|
+
images = [
|
|
127
|
+
"./path/to/image1.jpg",
|
|
128
|
+
"./path/to/image2.jpg",
|
|
129
|
+
]
|
|
130
|
+
|
|
131
|
+
model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
|
|
132
|
+
embeddings = list(embedding_model.embed(images))
|
|
133
|
+
|
|
134
|
+
# [
|
|
135
|
+
# array([-0.1115, 0.0097, 0.0052, 0.0195, ...], dtype=float32),
|
|
136
|
+
# array([-0.1019, 0.0635, -0.0332, 0.0522, ...], dtype=float32)
|
|
137
|
+
# ]
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
## ⚡️ FastEmbed on a GPU
|
|
142
|
+
|
|
143
|
+
FastEmbed supports running on GPU devices.
|
|
144
|
+
It requires installation of the `fastembed-gpu` package.
|
|
55
145
|
|
|
56
146
|
```bash
|
|
57
147
|
pip install fastembed-gpu
|
|
58
|
-
```
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
|
|
59
151
|
|
|
60
152
|
```python
|
|
61
153
|
from fastembed import TextEmbedding
|
|
62
154
|
|
|
63
|
-
embedding_model = TextEmbedding(
|
|
155
|
+
embedding_model = TextEmbedding(
|
|
156
|
+
model_name="BAAI/bge-small-en-v1.5",
|
|
157
|
+
providers=["CUDAExecutionProvider"]
|
|
158
|
+
)
|
|
64
159
|
print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
|
|
65
160
|
|
|
66
161
|
```
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
import importlib.metadata
|
|
2
|
+
|
|
3
|
+
from fastembed.image import ImageEmbedding
|
|
4
|
+
from fastembed.late_interaction import LateInteractionTextEmbedding
|
|
5
|
+
from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
|
|
6
|
+
from fastembed.text import TextEmbedding
|
|
7
|
+
|
|
8
|
+
try:
|
|
9
|
+
version = importlib.metadata.version("fastembed")
|
|
10
|
+
except importlib.metadata.PackageNotFoundError as _:
|
|
11
|
+
version = importlib.metadata.version("fastembed-gpu")
|
|
12
|
+
|
|
13
|
+
__version__ = version
|
|
14
|
+
__all__ = [
|
|
15
|
+
"TextEmbedding",
|
|
16
|
+
"SparseTextEmbedding",
|
|
17
|
+
"SparseEmbedding",
|
|
18
|
+
"ImageEmbedding",
|
|
19
|
+
"LateInteractionTextEmbedding",
|
|
20
|
+
]
|
|
@@ -2,13 +2,13 @@ import os
|
|
|
2
2
|
import shutil
|
|
3
3
|
import tarfile
|
|
4
4
|
from pathlib import Path
|
|
5
|
-
from typing import
|
|
5
|
+
from typing import Any, Dict, List, Optional
|
|
6
6
|
|
|
7
7
|
import requests
|
|
8
8
|
from huggingface_hub import snapshot_download
|
|
9
9
|
from huggingface_hub.utils import RepositoryNotFoundError
|
|
10
|
-
from tqdm import tqdm
|
|
11
10
|
from loguru import logger
|
|
11
|
+
from tqdm import tqdm
|
|
12
12
|
|
|
13
13
|
|
|
14
14
|
class ModelManagement:
|
|
@@ -42,7 +42,9 @@ class ModelManagement:
|
|
|
42
42
|
raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
|
|
43
43
|
|
|
44
44
|
@classmethod
|
|
45
|
-
def download_file_from_gcs(
|
|
45
|
+
def download_file_from_gcs(
|
|
46
|
+
cls, url: str, output_path: str, show_progress: bool = True
|
|
47
|
+
) -> str:
|
|
46
48
|
"""
|
|
47
49
|
Downloads a file from Google Cloud Storage.
|
|
48
50
|
|
|
@@ -71,12 +73,17 @@ class ModelManagement:
|
|
|
71
73
|
|
|
72
74
|
# Warn if the total size is zero
|
|
73
75
|
if total_size_in_bytes == 0:
|
|
74
|
-
print(
|
|
76
|
+
print(
|
|
77
|
+
f"Warning: Content-length header is missing or zero in the response from {url}."
|
|
78
|
+
)
|
|
75
79
|
|
|
76
80
|
show_progress = total_size_in_bytes and show_progress
|
|
77
81
|
|
|
78
82
|
with tqdm(
|
|
79
|
-
total=total_size_in_bytes,
|
|
83
|
+
total=total_size_in_bytes,
|
|
84
|
+
unit="iB",
|
|
85
|
+
unit_scale=True,
|
|
86
|
+
disable=not show_progress,
|
|
80
87
|
) as progress_bar:
|
|
81
88
|
with open(output_path, "wb") as file:
|
|
82
89
|
for chunk in response.iter_content(chunk_size=1024):
|
|
@@ -108,6 +115,7 @@ class ModelManagement:
|
|
|
108
115
|
"tokenizer.json",
|
|
109
116
|
"tokenizer_config.json",
|
|
110
117
|
"special_tokens_map.json",
|
|
118
|
+
"preprocessor_config.json",
|
|
111
119
|
]
|
|
112
120
|
if extra_patterns is not None:
|
|
113
121
|
allow_patterns.extend(extra_patterns)
|
|
@@ -155,7 +163,9 @@ class ModelManagement:
|
|
|
155
163
|
return cache_dir
|
|
156
164
|
|
|
157
165
|
@classmethod
|
|
158
|
-
def retrieve_model_gcs(
|
|
166
|
+
def retrieve_model_gcs(
|
|
167
|
+
cls, model_name: str, source_url: str, cache_dir: str
|
|
168
|
+
) -> Path:
|
|
159
169
|
fast_model_name = f"fast-{model_name.split('/')[-1]}"
|
|
160
170
|
|
|
161
171
|
cache_tmp_dir = Path(cache_dir) / "tmp"
|
|
@@ -181,8 +191,12 @@ class ModelManagement:
|
|
|
181
191
|
output_path=str(model_tar_gz),
|
|
182
192
|
)
|
|
183
193
|
|
|
184
|
-
cls.decompress_to_cache(
|
|
185
|
-
|
|
194
|
+
cls.decompress_to_cache(
|
|
195
|
+
targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir)
|
|
196
|
+
)
|
|
197
|
+
assert (
|
|
198
|
+
model_tmp_dir.exists()
|
|
199
|
+
), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
|
|
186
200
|
|
|
187
201
|
model_tar_gz.unlink()
|
|
188
202
|
# Rename from tmp to final name is atomic
|
|
@@ -0,0 +1,123 @@
|
|
|
1
|
+
import warnings
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
from typing import (
|
|
5
|
+
Any,
|
|
6
|
+
Dict,
|
|
7
|
+
Generic,
|
|
8
|
+
Iterable,
|
|
9
|
+
Optional,
|
|
10
|
+
Sequence,
|
|
11
|
+
Tuple,
|
|
12
|
+
Type,
|
|
13
|
+
TypeVar,
|
|
14
|
+
)
|
|
15
|
+
|
|
16
|
+
import numpy as np
|
|
17
|
+
import onnxruntime as ort
|
|
18
|
+
|
|
19
|
+
from fastembed.common.types import OnnxProvider
|
|
20
|
+
from fastembed.parallel_processor import Worker
|
|
21
|
+
|
|
22
|
+
# Holds type of the embedding result
|
|
23
|
+
T = TypeVar("T")
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@dataclass
|
|
27
|
+
class OnnxOutputContext:
|
|
28
|
+
model_output: np.ndarray
|
|
29
|
+
attention_mask: Optional[np.ndarray] = None
|
|
30
|
+
input_ids: Optional[np.ndarray] = None
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class OnnxModel(Generic[T]):
|
|
34
|
+
@classmethod
|
|
35
|
+
def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
|
|
36
|
+
raise NotImplementedError("Subclasses must implement this method")
|
|
37
|
+
|
|
38
|
+
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
|
|
39
|
+
raise NotImplementedError("Subclasses must implement this method")
|
|
40
|
+
|
|
41
|
+
def __init__(self) -> None:
|
|
42
|
+
self.model = None
|
|
43
|
+
self.tokenizer = None
|
|
44
|
+
|
|
45
|
+
def _preprocess_onnx_input(
|
|
46
|
+
self, onnx_input: Dict[str, np.ndarray], **kwargs
|
|
47
|
+
) -> Dict[str, np.ndarray]:
|
|
48
|
+
"""
|
|
49
|
+
Preprocess the onnx input.
|
|
50
|
+
"""
|
|
51
|
+
return onnx_input
|
|
52
|
+
|
|
53
|
+
def load_onnx_model(
|
|
54
|
+
self,
|
|
55
|
+
model_dir: Path,
|
|
56
|
+
model_file: str,
|
|
57
|
+
threads: Optional[int],
|
|
58
|
+
providers: Optional[Sequence[OnnxProvider]] = None,
|
|
59
|
+
) -> None:
|
|
60
|
+
model_path = model_dir / model_file
|
|
61
|
+
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
|
|
62
|
+
|
|
63
|
+
onnx_providers = (
|
|
64
|
+
["CPUExecutionProvider"] if providers is None else list(providers)
|
|
65
|
+
)
|
|
66
|
+
available_providers = ort.get_available_providers()
|
|
67
|
+
requested_provider_names = []
|
|
68
|
+
for provider in onnx_providers:
|
|
69
|
+
# check providers available
|
|
70
|
+
provider_name = provider if isinstance(provider, str) else provider[0]
|
|
71
|
+
requested_provider_names.append(provider_name)
|
|
72
|
+
if provider_name not in available_providers:
|
|
73
|
+
raise ValueError(
|
|
74
|
+
f"Provider {provider_name} is not available. Available providers: {available_providers}"
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
so = ort.SessionOptions()
|
|
78
|
+
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
|
79
|
+
|
|
80
|
+
if threads is not None:
|
|
81
|
+
so.intra_op_num_threads = threads
|
|
82
|
+
so.inter_op_num_threads = threads
|
|
83
|
+
|
|
84
|
+
self.model = ort.InferenceSession(
|
|
85
|
+
str(model_path), providers=onnx_providers, sess_options=so
|
|
86
|
+
)
|
|
87
|
+
if "CUDAExecutionProvider" in requested_provider_names:
|
|
88
|
+
current_providers = self.model.get_providers()
|
|
89
|
+
if "CUDAExecutionProvider" not in current_providers:
|
|
90
|
+
warnings.warn(
|
|
91
|
+
f"Attempt to set CUDAExecutionProvider failed. Current providers: {current_providers}."
|
|
92
|
+
"If you are using CUDA 12.x, install onnxruntime-gpu via "
|
|
93
|
+
"`pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/`",
|
|
94
|
+
RuntimeWarning,
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
|
|
98
|
+
raise NotImplementedError("Subclasses must implement this method")
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
class EmbeddingWorker(Worker):
|
|
102
|
+
def init_embedding(
|
|
103
|
+
self,
|
|
104
|
+
model_name: str,
|
|
105
|
+
cache_dir: str,
|
|
106
|
+
**kwargs,
|
|
107
|
+
) -> OnnxModel:
|
|
108
|
+
raise NotImplementedError()
|
|
109
|
+
|
|
110
|
+
def __init__(
|
|
111
|
+
self,
|
|
112
|
+
model_name: str,
|
|
113
|
+
cache_dir: str,
|
|
114
|
+
**kwargs,
|
|
115
|
+
):
|
|
116
|
+
self.model = self.init_embedding(model_name, cache_dir, **kwargs)
|
|
117
|
+
|
|
118
|
+
@classmethod
|
|
119
|
+
def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
|
|
120
|
+
return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
|
|
121
|
+
|
|
122
|
+
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
|
|
123
|
+
raise NotImplementedError("Subclasses must implement this method")
|
|
@@ -1,11 +1,24 @@
|
|
|
1
1
|
import json
|
|
2
2
|
from pathlib import Path
|
|
3
|
+
from typing import Tuple
|
|
3
4
|
|
|
4
|
-
import
|
|
5
|
-
from tokenizers import Tokenizer, AddedToken
|
|
5
|
+
from tokenizers import AddedToken, Tokenizer
|
|
6
6
|
|
|
7
|
+
from fastembed.image.transform.operators import Compose
|
|
7
8
|
|
|
8
|
-
|
|
9
|
+
|
|
10
|
+
def load_special_tokens(model_dir: Path) -> dict:
|
|
11
|
+
tokens_map_path = model_dir / "special_tokens_map.json"
|
|
12
|
+
if not tokens_map_path.exists():
|
|
13
|
+
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
|
|
14
|
+
|
|
15
|
+
with open(str(tokens_map_path)) as tokens_map_file:
|
|
16
|
+
tokens_map = json.load(tokens_map_file)
|
|
17
|
+
|
|
18
|
+
return tokens_map
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, dict]:
|
|
9
22
|
config_path = model_dir / "config.json"
|
|
10
23
|
if not config_path.exists():
|
|
11
24
|
raise ValueError(f"Could not find config.json in {model_dir}")
|
|
@@ -18,21 +31,18 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
|
|
|
18
31
|
if not tokenizer_config_path.exists():
|
|
19
32
|
raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
|
|
20
33
|
|
|
21
|
-
tokens_map_path = model_dir / "special_tokens_map.json"
|
|
22
|
-
if not tokens_map_path.exists():
|
|
23
|
-
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
|
|
24
|
-
|
|
25
34
|
with open(str(config_path)) as config_file:
|
|
26
35
|
config = json.load(config_file)
|
|
27
36
|
|
|
28
37
|
with open(str(tokenizer_config_path)) as tokenizer_config_file:
|
|
29
38
|
tokenizer_config = json.load(tokenizer_config_file)
|
|
30
39
|
|
|
31
|
-
|
|
32
|
-
tokens_map = json.load(tokens_map_file)
|
|
40
|
+
tokens_map = load_special_tokens(model_dir)
|
|
33
41
|
|
|
34
42
|
tokenizer = Tokenizer.from_file(str(tokenizer_path))
|
|
35
|
-
tokenizer.enable_truncation(
|
|
43
|
+
tokenizer.enable_truncation(
|
|
44
|
+
max_length=min(tokenizer_config["model_max_length"], max_length)
|
|
45
|
+
)
|
|
36
46
|
tokenizer.enable_padding(
|
|
37
47
|
pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
|
|
38
48
|
)
|
|
@@ -43,12 +53,24 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
|
|
|
43
53
|
elif isinstance(token, dict):
|
|
44
54
|
tokenizer.add_special_tokens([AddedToken(**token)])
|
|
45
55
|
|
|
46
|
-
|
|
56
|
+
special_token_to_id = {}
|
|
57
|
+
|
|
58
|
+
for token in tokens_map.values():
|
|
59
|
+
if isinstance(token, str):
|
|
60
|
+
special_token_to_id[token] = tokenizer.token_to_id(token)
|
|
61
|
+
elif isinstance(token, dict):
|
|
62
|
+
token_str = token.get("content", "")
|
|
63
|
+
special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
|
|
64
|
+
|
|
65
|
+
return tokenizer, special_token_to_id
|
|
66
|
+
|
|
47
67
|
|
|
68
|
+
def load_preprocessor(model_dir: Path) -> Compose:
|
|
69
|
+
preprocessor_config_path = model_dir / "preprocessor_config.json"
|
|
70
|
+
if not preprocessor_config_path.exists():
|
|
71
|
+
raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
|
|
48
72
|
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
normalized_array = input_array / norm
|
|
54
|
-
return normalized_array
|
|
73
|
+
with open(str(preprocessor_config_path)) as preprocessor_config_file:
|
|
74
|
+
preprocessor_config = json.load(preprocessor_config_file)
|
|
75
|
+
transforms = Compose.from_config(preprocessor_config)
|
|
76
|
+
return transforms
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import sys
|
|
3
|
+
from typing import Any, Dict, Iterable, Tuple, Union
|
|
4
|
+
|
|
5
|
+
if sys.version_info >= (3, 10):
|
|
6
|
+
from typing import TypeAlias
|
|
7
|
+
else:
|
|
8
|
+
from typing_extensions import TypeAlias
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
PathInput: TypeAlias = Union[str, os.PathLike]
|
|
12
|
+
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
|
|
13
|
+
|
|
14
|
+
OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]
|
|
@@ -2,7 +2,17 @@ import os
|
|
|
2
2
|
import tempfile
|
|
3
3
|
from itertools import islice
|
|
4
4
|
from pathlib import Path
|
|
5
|
-
from typing import
|
|
5
|
+
from typing import Generator, Iterable, Optional, Union
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
|
|
11
|
+
# Calculate the Lp norm along the specified dimension
|
|
12
|
+
norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
|
|
13
|
+
norm = np.maximum(norm, eps) # Avoid division by zero
|
|
14
|
+
normalized_array = input_array / norm
|
|
15
|
+
return normalized_array
|
|
6
16
|
|
|
7
17
|
|
|
8
18
|
def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
|