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.
Files changed (45) hide show
  1. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/PKG-INFO +114 -15
  2. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/README.md +103 -8
  3. fastembed_gpu-0.3.1/fastembed/__init__.py +20 -0
  4. fastembed_gpu-0.3.1/fastembed/common/__init__.py +3 -0
  5. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/common/model_management.py +22 -8
  6. fastembed_gpu-0.3.1/fastembed/common/onnx_model.py +123 -0
  7. fastembed_gpu-0.2.7/fastembed/common/models.py → fastembed_gpu-0.3.1/fastembed/common/preprocessor_utils.py +39 -17
  8. fastembed_gpu-0.3.1/fastembed/common/types.py +14 -0
  9. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/common/utils.py +11 -1
  10. fastembed_gpu-0.3.1/fastembed/image/__init__.py +3 -0
  11. fastembed_gpu-0.3.1/fastembed/image/image_embedding.py +94 -0
  12. fastembed_gpu-0.3.1/fastembed/image/image_embedding_base.py +44 -0
  13. fastembed_gpu-0.3.1/fastembed/image/onnx_embedding.py +131 -0
  14. fastembed_gpu-0.3.1/fastembed/image/onnx_image_model.py +108 -0
  15. fastembed_gpu-0.3.1/fastembed/image/transform/functional.py +124 -0
  16. fastembed_gpu-0.3.1/fastembed/image/transform/operators.py +198 -0
  17. fastembed_gpu-0.3.1/fastembed/late_interaction/__init__.py +5 -0
  18. fastembed_gpu-0.3.1/fastembed/late_interaction/colbert.py +194 -0
  19. fastembed_gpu-0.3.1/fastembed/late_interaction/late_interaction_embedding_base.py +62 -0
  20. fastembed_gpu-0.3.1/fastembed/late_interaction/late_interaction_text_embedding.py +109 -0
  21. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/parallel_processor.py +6 -2
  22. fastembed_gpu-0.3.1/fastembed/sparse/bm25.py +284 -0
  23. fastembed_gpu-0.3.1/fastembed/sparse/bm42.py +292 -0
  24. fastembed_gpu-0.3.1/fastembed/sparse/sparse_embedding_base.py +85 -0
  25. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/sparse/sparse_text_embedding.py +31 -7
  26. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/sparse/splade_pp.py +18 -21
  27. fastembed_gpu-0.3.1/fastembed/sparse/utils/tokenizer.py +113 -0
  28. fastembed_gpu-0.3.1/fastembed/text/clip_embedding.py +49 -0
  29. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/e5_onnx_embedding.py +10 -8
  30. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/jina_onnx_embedding.py +21 -12
  31. fastembed_gpu-0.3.1/fastembed/text/mini_lm_embedding.py +58 -0
  32. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/onnx_embedding.py +18 -27
  33. fastembed_gpu-0.3.1/fastembed/text/onnx_text_model.py +126 -0
  34. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/text_embedding.py +14 -3
  35. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/text/text_embedding_base.py +3 -1
  36. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/pyproject.toml +11 -7
  37. fastembed_gpu-0.2.7/fastembed/__init__.py +0 -12
  38. fastembed_gpu-0.2.7/fastembed/common/__init__.py +0 -3
  39. fastembed_gpu-0.2.7/fastembed/common/onnx_model.py +0 -165
  40. fastembed_gpu-0.2.7/fastembed/image/__init__.py +0 -0
  41. fastembed_gpu-0.2.7/fastembed/sparse/sparse_embedding_base.py +0 -44
  42. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/LICENSE +0 -0
  43. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/embedding.py +0 -0
  44. {fastembed_gpu-0.2.7 → fastembed_gpu-0.3.1}/fastembed/sparse/__init__.py +0 -0
  45. {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.2.7
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: NirantK
9
- Author-email: nirant.bits@gmail.com
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: huggingface-hub (>=0.20,<0.21)
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: numpy (>=1.21) ; python_version < "3.12"
21
- Requires-Dist: numpy (>=1.26) ; python_version >= "3.12"
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: tokenizers (>=0.15,<0.16)
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
- ### ⚡️ With GPU
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
- ### ⚡️ FastEmbed on a GPU
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
- FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
82
- Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
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(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
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
- ### ⚡️ With GPU
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
- ### ⚡️ FastEmbed on a GPU
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
- FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
54
- Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
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(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
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
+ ]
@@ -0,0 +1,3 @@
1
+ from fastembed.common.types import ImageInput, OnnxProvider, PathInput
2
+
3
+ __all__ = ["OnnxProvider", "ImageInput", "PathInput"]
@@ -2,13 +2,13 @@ import os
2
2
  import shutil
3
3
  import tarfile
4
4
  from pathlib import Path
5
- from typing import List, Optional, Dict, Any
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(cls, url: str, output_path: str, show_progress: bool = True) -> str:
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(f"Warning: Content-length header is missing or zero in the response from {url}.")
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, unit="iB", unit_scale=True, disable=not show_progress
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(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
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(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
185
- assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
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 numpy as np
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
- def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
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
- with open(str(tokens_map_path)) as tokens_map_file:
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(max_length=min(tokenizer_config["model_max_length"], max_length))
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
- return tokenizer
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
- def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
50
- # Calculate the Lp norm along the specified dimension
51
- norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
52
- norm = np.maximum(norm, eps) # Avoid division by zero
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 Union, Iterable, Generator, Optional
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:
@@ -0,0 +1,3 @@
1
+ from fastembed.image.image_embedding import ImageEmbedding
2
+
3
+ __all__ = ["ImageEmbedding"]