mirror of
https://github.com/qdrant/fastembed.git
synced 2026-09-23 06:27:51 -05:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Nirant Kasliwal
|
6bae4c57d3 |
@@ -1,6 +1,6 @@
|
||||
# ⚡️ What is FastEmbed?
|
||||
|
||||
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.
|
||||
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, the top model 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/).
|
||||
|
||||
@@ -23,7 +23,7 @@ To install the FastEmbed library, pip works:
|
||||
pip install fastembed
|
||||
```
|
||||
|
||||
## 📖 Quickstart
|
||||
## 📖 Usage
|
||||
|
||||
```python
|
||||
from fastembed import TextEmbedding
|
||||
@@ -48,14 +48,16 @@ Installation with Qdrant Client in Python:
|
||||
pip install qdrant-client[fastembed]
|
||||
```
|
||||
|
||||
You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
Might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
|
||||
```python
|
||||
from qdrant_client import QdrantClient
|
||||
|
||||
# Initialize the client
|
||||
client = QdrantClient("localhost", port=6333) # For production
|
||||
# client = QdrantClient(":memory:") # For small experiments
|
||||
# OR if you just want to try it out quickly:
|
||||
# client = QdrantClient(":memory:")
|
||||
# client = QdrantClient(path="path/to/db")
|
||||
|
||||
# Prepare your documents, metadata, and IDs
|
||||
docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
|
||||
@@ -65,12 +67,7 @@ metadata = [
|
||||
]
|
||||
ids = [42, 2]
|
||||
|
||||
# If you want to change the model:
|
||||
# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
|
||||
# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
|
||||
|
||||
# Use the new add() instead of upsert()
|
||||
# This internally calls embed() of the configured embedding model
|
||||
# Use the new add method
|
||||
client.add(
|
||||
collection_name="demo_collection",
|
||||
documents=docs,
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
@@ -110,30 +110,6 @@
|
||||
" </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.54</td>\n",
|
||||
" <td>{'hf': 'nomic-ai/nomic-embed-text-v1'}</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.54</td>\n",
|
||||
" <td>{'hf': 'nomic-ai/nomic-embed-text-v1.5'}</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>10</th>\n",
|
||||
" <td>thenlper/gte-large</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Large general text embeddings model</td>\n",
|
||||
" <td>1.34</td>\n",
|
||||
" <td>{'hf': 'qdrant/gte-large-onnx'}</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>11</th>\n",
|
||||
" <td>intfloat/multilingual-e5-large</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Multilingual model, e5-large. Recommend using this model for non-English languages</td>\n",
|
||||
@@ -141,7 +117,7 @@
|
||||
" <td>{'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'}</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>12</th>\n",
|
||||
" <th>9</th>\n",
|
||||
" <td>sentence-transformers/paraphrase-multilingual-mpnet-base-v2</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Sentence-transformers model for tasks like clustering or semantic search</td>\n",
|
||||
@@ -149,15 +125,7 @@
|
||||
" <td>{'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>13</th>\n",
|
||||
" <td>sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2</td>\n",
|
||||
" <td>384</td>\n",
|
||||
" <td>Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2</td>\n",
|
||||
" <td>0.46</td>\n",
|
||||
" <td>{'hf': 'qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q'}</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>14</th>\n",
|
||||
" <th>10</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-en</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>English embedding model supporting 8192 sequence length</td>\n",
|
||||
@@ -165,7 +133,7 @@
|
||||
" <td>{'hf': 'xenova/jina-embeddings-v2-base-en'}</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>15</th>\n",
|
||||
" <th>11</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-small-en</td>\n",
|
||||
" <td>512</td>\n",
|
||||
" <td>English embedding model supporting 8192 sequence length</td>\n",
|
||||
@@ -186,14 +154,10 @@
|
||||
"5 BAAI/bge-small-en-v1.5 384 \n",
|
||||
"6 BAAI/bge-small-zh-v1.5 512 \n",
|
||||
"7 sentence-transformers/all-MiniLM-L6-v2 384 \n",
|
||||
"8 nomic-ai/nomic-embed-text-v1 768 \n",
|
||||
"9 nomic-ai/nomic-embed-text-v1.5 768 \n",
|
||||
"10 thenlper/gte-large 1024 \n",
|
||||
"11 intfloat/multilingual-e5-large 1024 \n",
|
||||
"12 sentence-transformers/paraphrase-multilingual-mpnet-base-v2 768 \n",
|
||||
"13 sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 384 \n",
|
||||
"14 jinaai/jina-embeddings-v2-base-en 768 \n",
|
||||
"15 jinaai/jina-embeddings-v2-small-en 512 \n",
|
||||
"8 intfloat/multilingual-e5-large 1024 \n",
|
||||
"9 sentence-transformers/paraphrase-multilingual-mpnet-base-v2 768 \n",
|
||||
"10 jinaai/jina-embeddings-v2-base-en 768 \n",
|
||||
"11 jinaai/jina-embeddings-v2-small-en 512 \n",
|
||||
"\n",
|
||||
" description \\\n",
|
||||
"0 Base English model \n",
|
||||
@@ -204,14 +168,10 @@
|
||||
"5 Fast and Default English model \n",
|
||||
"6 Fast and recommended Chinese model \n",
|
||||
"7 Sentence Transformer model, MiniLM-L6-v2 \n",
|
||||
"8 8192 context length english model \n",
|
||||
"9 8192 context length english model \n",
|
||||
"10 Large general text embeddings model \n",
|
||||
"11 Multilingual model, e5-large. Recommend using this model for non-English languages \n",
|
||||
"12 Sentence-transformers model for tasks like clustering or semantic search \n",
|
||||
"13 Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2 \n",
|
||||
"14 English embedding model supporting 8192 sequence length \n",
|
||||
"15 English embedding model supporting 8192 sequence length \n",
|
||||
"8 Multilingual model, e5-large. Recommend using this model for non-English languages \n",
|
||||
"9 Sentence-transformers model for tasks like clustering or semantic search \n",
|
||||
"10 English embedding model supporting 8192 sequence length \n",
|
||||
"11 English embedding model supporting 8192 sequence length \n",
|
||||
"\n",
|
||||
" size_in_GB \\\n",
|
||||
"0 0.50 \n",
|
||||
@@ -222,14 +182,10 @@
|
||||
"5 0.13 \n",
|
||||
"6 0.10 \n",
|
||||
"7 0.09 \n",
|
||||
"8 0.54 \n",
|
||||
"9 0.54 \n",
|
||||
"10 1.34 \n",
|
||||
"11 2.24 \n",
|
||||
"12 1.11 \n",
|
||||
"13 0.46 \n",
|
||||
"14 0.55 \n",
|
||||
"15 0.13 \n",
|
||||
"8 2.24 \n",
|
||||
"9 1.11 \n",
|
||||
"10 0.55 \n",
|
||||
"11 0.13 \n",
|
||||
"\n",
|
||||
" sources \n",
|
||||
"0 {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz'} \n",
|
||||
@@ -240,17 +196,13 @@
|
||||
"5 {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz', 'hf': 'qdrant/bge-small-en-v1.5-onnx-q'} \n",
|
||||
"6 {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz'} \n",
|
||||
"7 {'url': 'https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz', 'hf': 'qdrant/all-MiniLM-L6-v2-onnx'} \n",
|
||||
"8 {'hf': 'nomic-ai/nomic-embed-text-v1'} \n",
|
||||
"9 {'hf': 'nomic-ai/nomic-embed-text-v1.5'} \n",
|
||||
"10 {'hf': 'qdrant/gte-large-onnx'} \n",
|
||||
"11 {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'} \n",
|
||||
"12 {'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'} \n",
|
||||
"13 {'hf': 'qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q'} \n",
|
||||
"14 {'hf': 'xenova/jina-embeddings-v2-base-en'} \n",
|
||||
"15 {'hf': 'xenova/jina-embeddings-v2-small-en'} "
|
||||
"8 {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'} \n",
|
||||
"9 {'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'} \n",
|
||||
"10 {'hf': 'xenova/jina-embeddings-v2-base-en'} \n",
|
||||
"11 {'hf': 'xenova/jina-embeddings-v2-small-en'} "
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
@@ -280,7 +232,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.13"
|
||||
"version": "3.11.5"
|
||||
},
|
||||
"orig_nbformat": 4
|
||||
},
|
||||
|
||||
@@ -1,473 +0,0 @@
|
||||
"""
|
||||
This script is used to convert HuggingFace models to ONNX format and optionally quantize the model using dynamic quantization.
|
||||
This is courtesy of Joshua aka @Xenova
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Set
|
||||
|
||||
import onnx
|
||||
from onnxruntime.quantization import QuantType, quantize_dynamic
|
||||
from optimum.exporters.onnx import export_models, main_export
|
||||
from optimum.exporters.tasks import TasksManager
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoConfig, AutoTokenizer, HfArgumentParser
|
||||
|
||||
DEFAULT_QUANTIZE_PARAMS = {
|
||||
"per_channel": True,
|
||||
"reduce_range": True,
|
||||
}
|
||||
|
||||
MODEL_SPECIFIC_QUANTIZE_PARAMS = {
|
||||
# Decoder-only models
|
||||
"codegen": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt2": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt_bigcode": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gptj": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt-neo": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"gpt-neox": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"mpt": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"bloom": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"llama": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"opt": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"mistral": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"falcon": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"phi": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"qwen2": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
# Encoder-decoder models
|
||||
"whisper": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
"vision-encoder-decoder": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
# Encoder-only models
|
||||
"owlv2": {
|
||||
"per_channel": False,
|
||||
"reduce_range": False,
|
||||
},
|
||||
}
|
||||
|
||||
MODELS_WITHOUT_TOKENIZERS = [
|
||||
"wav2vec2",
|
||||
"wav2vec2-bert",
|
||||
"wavlm",
|
||||
"hubert",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversionArguments:
|
||||
"""
|
||||
Arguments used for converting HuggingFace models to onnx.
|
||||
"""
|
||||
|
||||
model_id: str = field(metadata={"help": "Model identifier"})
|
||||
tokenizer_id: str = field(default=None, metadata={"help": "Tokenizer identifier (if different to `model_id`)"})
|
||||
quantize: bool = field(default=False, metadata={"help": "Whether to quantize the model."})
|
||||
output_parent_dir: str = field(
|
||||
default="./models/", metadata={"help": "Path where the converted model will be saved to."}
|
||||
)
|
||||
|
||||
task: Optional[str] = field(
|
||||
default="auto",
|
||||
metadata={
|
||||
"help": (
|
||||
"The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
|
||||
f" {str(TasksManager.get_all_tasks())}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
opset: int = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": (
|
||||
"If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
device: str = field(default="cpu", metadata={"help": "The device to use to do the export."})
|
||||
skip_validation: bool = field(default=False, metadata={"help": "Whether to skip validation of the converted model"})
|
||||
|
||||
per_channel: bool = field(default=None, metadata={"help": "Whether to quantize weights per channel"})
|
||||
reduce_range: bool = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "Whether to quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode"
|
||||
},
|
||||
)
|
||||
|
||||
output_attentions: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Whether to output attentions from the model. NOTE: This is only supported for whisper models right now."
|
||||
},
|
||||
)
|
||||
|
||||
split_modalities: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Whether to split multimodal models. NOTE: This is only supported for CLIP models right now."
|
||||
},
|
||||
)
|
||||
|
||||
trust_remote_code: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories"
|
||||
"you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository."
|
||||
},
|
||||
)
|
||||
|
||||
custom_onnx_configs: str = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users "
|
||||
"that desire a finer-grained control on the export."
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def get_operators(model: onnx.ModelProto) -> Set[str]:
|
||||
operators = set()
|
||||
|
||||
def traverse_graph(graph):
|
||||
for node in graph.node:
|
||||
operators.add(node.op_type)
|
||||
for attr in node.attribute:
|
||||
if attr.type == onnx.AttributeProto.GRAPH:
|
||||
subgraph = attr.g
|
||||
traverse_graph(subgraph)
|
||||
|
||||
traverse_graph(model.graph)
|
||||
return operators
|
||||
|
||||
|
||||
def quantize(model_names_or_paths, **quantize_kwargs):
|
||||
"""
|
||||
Quantize the weights of the model from float32 to int8 to allow very efficient inference on modern CPU
|
||||
|
||||
Uses unsigned ints for activation values, signed ints for weights, per
|
||||
https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
|
||||
it is faster on most CPU architectures
|
||||
Args:
|
||||
onnx_model_path: Path to location the exported ONNX model is stored
|
||||
Returns: The Path generated for the quantized
|
||||
"""
|
||||
|
||||
quantize_config = dict(**quantize_kwargs, per_model_config={})
|
||||
|
||||
for model in tqdm(model_names_or_paths, desc="Quantizing"):
|
||||
directory_path = os.path.dirname(model)
|
||||
file_name_without_extension = os.path.splitext(os.path.basename(model))[0]
|
||||
|
||||
# NOTE:
|
||||
# As of 2023/04/20, the current latest version of onnxruntime-web is 1.14.0, and does not support INT8 weights for Conv layers.
|
||||
# For this reason, we choose model weight types to ensure compatibility with onnxruntime-web.
|
||||
#
|
||||
# As per docs, signed weight type (QInt8) is faster on most CPUs, so, we use that unless the model contains a Conv layer.
|
||||
# For more information, see:
|
||||
# - https://github.com/microsoft/onnxruntime/issues/3130#issuecomment-1105200621
|
||||
# - https://github.com/microsoft/onnxruntime/issues/2339
|
||||
|
||||
loaded_model = onnx.load_model(model)
|
||||
op_types = get_operators(loaded_model)
|
||||
weight_type = QuantType.QUInt8 if "Conv" in op_types else QuantType.QInt8
|
||||
|
||||
quantize_dynamic(
|
||||
model_input=model,
|
||||
model_output=os.path.join(directory_path, f"{file_name_without_extension}_quantized.onnx"),
|
||||
weight_type=weight_type,
|
||||
# TODO allow user to specify these
|
||||
# op_types_to_quantize=['MatMul', 'Add', 'Conv'],
|
||||
extra_options=dict(EnableSubgraph=True),
|
||||
**quantize_kwargs,
|
||||
)
|
||||
|
||||
quantize_config["per_model_config"][file_name_without_extension] = dict(
|
||||
op_types=list(op_types),
|
||||
weight_type=str(weight_type),
|
||||
)
|
||||
|
||||
# Save quantization config
|
||||
with open(os.path.join(directory_path, "quantize_config.json"), "w") as fp:
|
||||
json.dump(quantize_config, fp, indent=4)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((ConversionArguments,))
|
||||
(conv_args,) = parser.parse_args_into_dataclasses()
|
||||
|
||||
model_id = conv_args.model_id
|
||||
tokenizer_id = conv_args.tokenizer_id or model_id
|
||||
|
||||
output_model_folder = os.path.join(conv_args.output_parent_dir, model_id)
|
||||
|
||||
# Create output folder
|
||||
os.makedirs(output_model_folder, exist_ok=True)
|
||||
|
||||
from_pretrained_kwargs = dict(
|
||||
trust_remote_code=conv_args.trust_remote_code,
|
||||
)
|
||||
|
||||
# Saving the model config
|
||||
config = AutoConfig.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
custom_kwargs = {}
|
||||
if conv_args.custom_onnx_configs is not None:
|
||||
if conv_args.task == "auto":
|
||||
raise Exception("`--task` must be set when exporting with `--custom_onnx_configs`")
|
||||
custom_onnx_configs = json.loads(conv_args.custom_onnx_configs)
|
||||
|
||||
for key in custom_onnx_configs:
|
||||
onnx_configs = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]["onnx"]
|
||||
mapping = onnx_configs[conv_args.task]
|
||||
custom_onnx_configs[key] = mapping.func(config, **mapping.keywords)
|
||||
|
||||
custom_kwargs["custom_onnx_configs"] = custom_onnx_configs
|
||||
|
||||
tokenizer = None
|
||||
try:
|
||||
# Load tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, **from_pretrained_kwargs)
|
||||
|
||||
# To avoid inserting all chat templates into tokenizers.js, we save the chat template
|
||||
# to the tokenizer_config.json file, and load it when the tokenizer is loaded.
|
||||
if getattr(tokenizer, "chat_template", None) is None and getattr(tokenizer, "use_default_system_prompt", False):
|
||||
# No chat template specified, and we use the default
|
||||
setattr(tokenizer, "chat_template", tokenizer.default_chat_template)
|
||||
|
||||
except KeyError:
|
||||
pass # No Tokenizer
|
||||
|
||||
except Exception as e:
|
||||
if config.model_type not in MODELS_WITHOUT_TOKENIZERS:
|
||||
raise e
|
||||
|
||||
core_export_kwargs = dict(
|
||||
opset=conv_args.opset,
|
||||
device=conv_args.device,
|
||||
trust_remote_code=conv_args.trust_remote_code,
|
||||
**custom_kwargs,
|
||||
)
|
||||
|
||||
export_kwargs = dict(
|
||||
model_name_or_path=model_id,
|
||||
output=output_model_folder,
|
||||
task=conv_args.task,
|
||||
do_validation=not conv_args.skip_validation,
|
||||
library_name="transformers",
|
||||
**core_export_kwargs,
|
||||
)
|
||||
|
||||
# Handle special cases
|
||||
if config.model_type == "marian":
|
||||
from .extra.marian import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(model_id, tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type == "esm":
|
||||
from .extra.esm import generate_fast_tokenizer
|
||||
|
||||
fast_tokenizer = generate_fast_tokenizer(tokenizer)
|
||||
fast_tokenizer.save(os.path.join(output_model_folder, "tokenizer.json"))
|
||||
|
||||
elif config.model_type == "whisper":
|
||||
if conv_args.output_attentions:
|
||||
from .extra.whisper import get_main_export_kwargs
|
||||
|
||||
export_kwargs.update(**get_main_export_kwargs(config, "automatic-speech-recognition"))
|
||||
|
||||
elif config.model_type in ("wav2vec2", "wav2vec2-bert", "hubert"):
|
||||
if tokenizer is not None:
|
||||
from .extra.wav2vec2 import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type == "vits":
|
||||
if tokenizer is not None:
|
||||
from .extra.vits import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type == "speecht5":
|
||||
# TODO allow user to specify vocoder path
|
||||
export_kwargs["model_kwargs"] = {"vocoder": "microsoft/speecht5_hifigan"}
|
||||
|
||||
if tokenizer is not None:
|
||||
from .extra.speecht5 import generate_tokenizer_json
|
||||
|
||||
tokenizer_json = generate_tokenizer_json(tokenizer)
|
||||
|
||||
with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp:
|
||||
json.dump(tokenizer_json, fp, indent=4)
|
||||
|
||||
elif config.model_type in ("owlvit", "owlv2"):
|
||||
# Override default batch size to 1, needed because non-maximum suppression is performed for exporting.
|
||||
# For more information, see https://github.com/huggingface/optimum/blob/e3b7efb1257c011db907ef40ab340e795cc5684c/optimum/exporters/onnx/model_configs.py#L1028-L1032
|
||||
export_kwargs["batch_size"] = 1
|
||||
|
||||
else:
|
||||
pass # TODO
|
||||
|
||||
# Step 1. convert huggingface model to onnx
|
||||
if not conv_args.split_modalities:
|
||||
main_export(**export_kwargs)
|
||||
else:
|
||||
custom_export_kwargs = dict(
|
||||
output_dir=output_model_folder,
|
||||
**core_export_kwargs,
|
||||
)
|
||||
|
||||
if config.model_type == "clip":
|
||||
# Handle special case for exporting text and vision models separately
|
||||
from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection
|
||||
|
||||
from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig
|
||||
|
||||
text_model = CLIPTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
export_models(
|
||||
models_and_onnx_configs={
|
||||
"text_model": (text_model, CLIPTextModelWithProjectionOnnxConfig(text_model.config)),
|
||||
"vision_model": (vision_model, CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)),
|
||||
},
|
||||
**custom_export_kwargs,
|
||||
)
|
||||
|
||||
elif config.model_type == "siglip":
|
||||
# Handle special case for exporting text and vision models separately
|
||||
from transformers.models.siglip import SiglipTextModel, SiglipVisionModel
|
||||
|
||||
from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig
|
||||
|
||||
text_model = SiglipTextModel.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
vision_model = SiglipVisionModel.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
export_models(
|
||||
models_and_onnx_configs={
|
||||
"text_model": (text_model, SiglipTextModelOnnxConfig(text_model.config)),
|
||||
"vision_model": (vision_model, SiglipVisionModelOnnxConfig(vision_model.config)),
|
||||
},
|
||||
**custom_export_kwargs,
|
||||
)
|
||||
|
||||
# TODO: Enable once https://github.com/huggingface/optimum/pull/1552 is merged
|
||||
# elif config.model_type == 'clap':
|
||||
# # Handle special case for exporting text and audio models separately
|
||||
# from .extra.clap import ClapTextModelWithProjectionOnnxConfig, ClapAudioModelWithProjectionOnnxConfig
|
||||
# from transformers.models.clap import ClapTextModelWithProjection, ClapAudioModelWithProjection
|
||||
|
||||
# text_model = ClapTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
# audio_model = ClapAudioModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
|
||||
# export_models(
|
||||
# models_and_onnx_configs={
|
||||
# "text_model": (text_model, ClapTextModelWithProjectionOnnxConfig(text_model.config)),
|
||||
# "audio_model": (audio_model, ClapAudioModelWithProjectionOnnxConfig(audio_model.config)),
|
||||
# },
|
||||
# **custom_export_kwargs,
|
||||
# )
|
||||
|
||||
else:
|
||||
raise Exception(f"Unable to export {config.model_type} model with `--split_modalities`.")
|
||||
|
||||
# Step 2. (optional, recommended) quantize the converted model for fast inference and to reduce model size.
|
||||
if conv_args.quantize:
|
||||
# Update quantize config with model specific defaults
|
||||
quantize_config = MODEL_SPECIFIC_QUANTIZE_PARAMS.get(config.model_type, DEFAULT_QUANTIZE_PARAMS)
|
||||
|
||||
# Update if user specified values
|
||||
if conv_args.per_channel is not None:
|
||||
quantize_config["per_channel"] = conv_args.per_channel
|
||||
|
||||
if conv_args.reduce_range is not None:
|
||||
quantize_config["reduce_range"] = conv_args.reduce_range
|
||||
|
||||
quantize(
|
||||
[
|
||||
os.path.join(output_model_folder, x)
|
||||
for x in os.listdir(output_model_folder)
|
||||
if x.endswith(".onnx") and not x.endswith("_quantized.onnx")
|
||||
],
|
||||
**quantize_config,
|
||||
)
|
||||
|
||||
# Step 3. Move .onnx files to the 'onnx' subfolder
|
||||
os.makedirs(os.path.join(output_model_folder, "onnx"), exist_ok=True)
|
||||
for file in os.listdir(output_model_folder):
|
||||
if file.endswith((".onnx", ".onnx_data")):
|
||||
shutil.move(os.path.join(output_model_folder, file), os.path.join(output_model_folder, "onnx", file))
|
||||
|
||||
# Step 4. Update the generation config if necessary
|
||||
if config.model_type == "whisper":
|
||||
from transformers import GenerationConfig
|
||||
|
||||
from .extra.whisper import get_alignment_heads
|
||||
|
||||
generation_config = GenerationConfig.from_pretrained(model_id, **from_pretrained_kwargs)
|
||||
generation_config.alignment_heads = get_alignment_heads(config)
|
||||
generation_config.save_pretrained(output_model_folder)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,56 +0,0 @@
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import click
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from optimum.onnxruntime import ORTModelForFeatureExtraction
|
||||
from optimum.pipelines import pipeline
|
||||
from torch import Tensor
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
|
||||
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
|
||||
last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
|
||||
return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
|
||||
|
||||
|
||||
def hf_embed(model_id: str, texts: List[str], tokenizer):
|
||||
# Tokenize the input texts
|
||||
model = AutoModel.from_pretrained(model_id)
|
||||
model.eval()
|
||||
encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
|
||||
|
||||
model_output = model(**encoded_input)
|
||||
sentence_embeddings = model_output[0][:, 0]
|
||||
sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
|
||||
return sentence_embeddings
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option("--model_id", help="model id from huggingface.co/models")
|
||||
@click.option("--model_dir", help="The person to greet.")
|
||||
def setup(model_id, model_dir):
|
||||
text = "This is a test sentence"
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
output_dir = Path(model_dir)
|
||||
model = ORTModelForFeatureExtraction.from_pretrained(output_dir)
|
||||
onnx_quant_embed = pipeline(
|
||||
"feature-extraction", model=model, accelerator="ort", tokenizer=tokenizer, return_tensors=True
|
||||
)
|
||||
quant_embeddings = onnx_quant_embed([text])
|
||||
quant_embeddings = F.normalize(quant_embeddings[0][:,0], p=2, dim=1)
|
||||
quant_embeddings = quant_embeddings.detach().numpy()
|
||||
print(quant_embeddings.shape)
|
||||
|
||||
torch_embeddings = hf_embed(model_id, texts=[text], tokenizer=tokenizer)
|
||||
torch_embeddings = F.normalize(torch_embeddings, p=2, dim=1)
|
||||
torch_embeddings = torch_embeddings.detach().numpy()
|
||||
print(torch_embeddings.shape)
|
||||
assert quant_embeddings.shape == torch_embeddings.shape
|
||||
print(np.allclose(quant_embeddings, torch_embeddings, atol=1e-5))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
setup()
|
||||
@@ -1,3 +1,3 @@
|
||||
from fastembed.text.text_embedding import TextEmbedding
|
||||
from .embedding import TextEmbedding
|
||||
|
||||
__all__ = ["TextEmbedding"]
|
||||
__all__ = [TextEmbedding]
|
||||
@@ -131,8 +131,7 @@ class ModelManagement:
|
||||
model_tmp_dir = cache_tmp_dir / fast_model_name
|
||||
model_dir = Path(cache_dir) / fast_model_name
|
||||
|
||||
# check if the model_dir and the model files are both present for macOS
|
||||
if model_dir.exists() and len(list(model_dir.glob("*"))) > 0:
|
||||
if model_dir.exists():
|
||||
return model_dir
|
||||
|
||||
if model_tmp_dir.exists():
|
||||
|
||||
@@ -33,7 +33,7 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
|
||||
|
||||
tokenizer = Tokenizer.from_file(str(tokenizer_path))
|
||||
tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
|
||||
tokenizer.enable_padding(pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"])
|
||||
tokenizer.enable_padding(pad_id=config["pad_token_id"], pad_token=tokenizer_config["pad_token"])
|
||||
|
||||
for token in tokens_map.values():
|
||||
if isinstance(token, str):
|
||||
|
||||
@@ -24,10 +24,10 @@ def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
|
||||
"""
|
||||
if cache_dir is None:
|
||||
default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
|
||||
cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
|
||||
cache_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
|
||||
else:
|
||||
cache_path = Path(cache_dir)
|
||||
cache_dir = Path(cache_dir)
|
||||
|
||||
cache_path.mkdir(parents=True, exist_ok=True)
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return cache_path
|
||||
return cache_dir
|
||||
|
||||
@@ -4,7 +4,7 @@ from loguru import logger
|
||||
|
||||
from fastembed.text.text_embedding import TextEmbedding
|
||||
|
||||
logger.warning("DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." "Use from fastembed import TextEmbedding instead.")
|
||||
logger.warning("DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." " Use TextEmbedding instead.")
|
||||
|
||||
DefaultEmbedding = TextEmbedding
|
||||
FlagEmbedding = TextEmbedding
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
[
|
||||
{
|
||||
"model": "BAAI/bge-base-en",
|
||||
"dim": 768,
|
||||
"description": "Base English model",
|
||||
"size_in_GB": 0.5,
|
||||
"hf_sources": [],
|
||||
"compressed_url_sources": [
|
||||
"https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz"
|
||||
]
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-base-en-v1.5",
|
||||
"dim": 768,
|
||||
"description": "Base English model, v1.5",
|
||||
"size_in_GB": 0.44,
|
||||
"hf_sources": [
|
||||
"qdrant/bge-base-en-v1.5-onnx-q"
|
||||
],
|
||||
"compressed_url_sources": [
|
||||
"https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz"
|
||||
]
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-large-en-v1.5",
|
||||
"dim": 1024,
|
||||
"description": "Large English model, v1.5",
|
||||
"size_in_GB": 1.34,
|
||||
"hf_sources": [
|
||||
"qdrant/bge-large-en-v1.5-onnx",
|
||||
"qdrant/bge-large-en-v1.5-onnx-q"
|
||||
],
|
||||
"compressed_url_sources": []
|
||||
},
|
||||
{
|
||||
"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,
|
||||
"description": "Fast and Default English model",
|
||||
"size_in_GB": 0.13,
|
||||
"hf_sources": [
|
||||
"qdrant/bge-small-en-v1.5-onnx-q"
|
||||
],
|
||||
"compressed_url_sources": [
|
||||
"https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz"
|
||||
]
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-small-zh-v1.5",
|
||||
"dim": 512,
|
||||
"description": "Fast and recommended Chinese model",
|
||||
"size_in_GB": 0.1,
|
||||
"hf_sources": [],
|
||||
"compressed_url_sources": [
|
||||
"https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz"
|
||||
]
|
||||
},
|
||||
{
|
||||
"model": "intfloat/multilingual-e5-large",
|
||||
"dim": 1024,
|
||||
"description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
|
||||
"size_in_GB": 2.24,
|
||||
"hf_sources": [
|
||||
"qdrant/multilingual-e5-large-onnx"
|
||||
],
|
||||
"compressed_url_sources": [
|
||||
"https://storage.googleapis.com/qdrant-fastembed/intfloat-multilingual-e5-large.tar.gz"
|
||||
]
|
||||
},
|
||||
{
|
||||
"model": "jinaai/jina-embeddings-v2-base-en",
|
||||
"dim": 768,
|
||||
"description": "English embedding model supporting 8192 sequence length",
|
||||
"size_in_GB": 0.55,
|
||||
"hf_sources": [
|
||||
"xenova/jina-embeddings-v2-base-en"
|
||||
],
|
||||
"compressed_url_sources": []
|
||||
},
|
||||
{
|
||||
"model": "jinaai/jina-embeddings-v2-small-en",
|
||||
"dim": 512,
|
||||
"description": " English embedding model supporting 8192 sequence length",
|
||||
"size_in_GB": 0.13,
|
||||
"hf_sources": [
|
||||
"xenova/jina-embeddings-v2-small-en"
|
||||
],
|
||||
"compressed_url_sources": []
|
||||
},
|
||||
{
|
||||
"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": "xenova/multilingual-e5-large",
|
||||
"dim": 1024,
|
||||
"description": "Multilingual model. Recommended for non-English languages",
|
||||
"size_in_GB": 2.24,
|
||||
"hf_sources": [
|
||||
"xenova/multilingual-e5-large"
|
||||
],
|
||||
"compressed_url_sources": []
|
||||
},
|
||||
{
|
||||
"model": "xenova/paraphrase-multilingual-mpnet-base-v2",
|
||||
"dim": 768,
|
||||
"description": "Sentence-transformers model for tasks like clustering or semantic search",
|
||||
"size_in_GB": 1.11,
|
||||
"hf_sources": [
|
||||
"xenova/paraphrase-multilingual-mpnet-base-v2"
|
||||
],
|
||||
"compressed_url_sources": []
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
Reference implementation of BGE_M3 model.
|
||||
https://github.com/FlagOpen/FlagEmbedding/blob/e23ff5e213350cbd4bb50883f6dbbecf6c267965/FlagEmbedding/BGE_M3/modeling.py#L340
|
||||
"""
|
||||
from typing import Dict
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class BGE_M3:
|
||||
|
||||
def dense_embedding_np(hidden_state, mask, sentence_pooling_method):
|
||||
if sentence_pooling_method == 'cls':
|
||||
return hidden_state[:, 0, :]
|
||||
elif sentence_pooling_method == 'mean':
|
||||
masked_hidden_state = hidden_state * mask[:, :, None]
|
||||
sum_embeddings = masked_hidden_state.sum(axis=1)
|
||||
token_counts = mask.sum(axis=1, keepdims=True)
|
||||
token_counts = np.where(token_counts == 0, 1, token_counts)
|
||||
mean_embeddings = sum_embeddings / token_counts
|
||||
return mean_embeddings
|
||||
|
||||
|
||||
def forward(self,
|
||||
text_input: Dict[str, Tensor] = None,
|
||||
return_dense: bool = True,
|
||||
return_sparse: bool = False,
|
||||
return_colbert: bool = False,
|
||||
return_sparse_embedding: bool = False):
|
||||
assert return_dense or return_sparse or return_colbert, 'Must choose one or more from `return_colbert`, `return_sparse`, `return_dense` to set `True`!'
|
||||
|
||||
last_hidden_state = self.model(**text_input, return_dict=True).last_hidden_state
|
||||
|
||||
output = {}
|
||||
if return_dense:
|
||||
dense_vecs = self.dense_embedding(last_hidden_state, text_input['attention_mask'])
|
||||
output['dense_vecs'] = dense_vecs
|
||||
if return_sparse:
|
||||
sparse_vecs = self.sparse_embedding(last_hidden_state, text_input['input_ids'],
|
||||
return_embedding=return_sparse_embedding)
|
||||
output['sparse_vecs'] = sparse_vecs
|
||||
if return_colbert:
|
||||
colbert_vecs = self.colbert_embedding(last_hidden_state, text_input['attention_mask'])
|
||||
output['colbert_vecs'] = colbert_vecs
|
||||
|
||||
if self.normlized:
|
||||
if 'dense_vecs' in output:
|
||||
output['dense_vecs'] = torch.nn.functional.normalize(output['dense_vecs'], dim=-1)
|
||||
if 'colbert_vecs' in output:
|
||||
output['colbert_vecs'] = torch.nn.functional.normalize(output['colbert_vecs'], dim=-1)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,17 @@
|
||||
import onnx
|
||||
from onnx import helper
|
||||
|
||||
# Load the original ONNX model
|
||||
model = onnx.load('path_to_your_original_model.onnx')
|
||||
|
||||
# Assuming 'colbert_linear' and 'sparse_linear' are the names of the intermediate nodes you're interested in
|
||||
# You need to find out the exact names by inspecting the model, e.g., using Netron
|
||||
|
||||
# Add these nodes as additional outputs to the model
|
||||
output_for_colbert_linear = helper.make_tensor_value_info('colbert_linear', onnx.TensorProto.FLOAT, [your_shape_here])
|
||||
output_for_sparse_linear = helper.make_tensor_value_info('sparse_linear', onnx.TensorProto.FLOAT, [your_shape_here])
|
||||
|
||||
model.graph.output.extend([output_for_colbert_linear, output_for_sparse_linear])
|
||||
|
||||
# Save the modified model
|
||||
onnx.save(model, 'path_to_your_modified_model.onnx')
|
||||
@@ -22,8 +22,8 @@ supported_multilingual_e5_models = [
|
||||
"size_in_GB": 1.11,
|
||||
"sources": {
|
||||
"hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@@ -51,8 +51,8 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
|
||||
|
||||
class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
|
||||
def init_embedding(
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
self,
|
||||
model_name: str,
|
||||
cache_dir: str,
|
||||
) -> E5OnnxEmbedding:
|
||||
return E5OnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import os
|
||||
from multiprocessing import get_all_start_methods
|
||||
from typing import List, Dict, Any, Optional, Tuple, Union, Iterable, Type
|
||||
from typing import List, Dict, Any, Tuple, Union, Iterable, Type
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
@@ -98,42 +98,6 @@ supported_onnx_models = [
|
||||
"hf": "qdrant/all-MiniLM-L6-v2-onnx",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
|
||||
"dim": 384,
|
||||
"description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
|
||||
"size_in_GB": 0.46,
|
||||
"sources": {
|
||||
"hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model": "nomic-ai/nomic-embed-text-v1",
|
||||
"dim": 768,
|
||||
"description": "8192 context length english model",
|
||||
"size_in_GB": 0.54,
|
||||
"sources": {
|
||||
"hf": "nomic-ai/nomic-embed-text-v1",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model": "nomic-ai/nomic-embed-text-v1.5",
|
||||
"dim": 768,
|
||||
"description": "8192 context length english model",
|
||||
"size_in_GB": 0.54,
|
||||
"sources": {
|
||||
"hf": "nomic-ai/nomic-embed-text-v1.5",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model": "thenlper/gte-large",
|
||||
"dim": 1024,
|
||||
"description": "Large general text embeddings model",
|
||||
"size_in_GB": 1.34,
|
||||
"sources": {
|
||||
"hf": "qdrant/gte-large-onnx",
|
||||
},
|
||||
},
|
||||
# {
|
||||
# "model": "sentence-transformers/all-MiniLM-L6-v2",
|
||||
# "dim": 384,
|
||||
@@ -185,8 +149,8 @@ class OnnxTextEmbedding(TextEmbeddingBase):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "BAAI/bge-small-en-v1.5",
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
cache_dir: str = None,
|
||||
threads: int = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
@@ -229,7 +193,7 @@ class OnnxTextEmbedding(TextEmbeddingBase):
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
parallel: int = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
|
||||
@@ -53,22 +53,24 @@ class TextEmbedding(TextEmbeddingBase):
|
||||
):
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
self.model = None
|
||||
for embedding in self.EMBEDDINGS_REGISTRY:
|
||||
supported_models = embedding.list_supported_models()
|
||||
if any(model_name == model["model"] for model in supported_models):
|
||||
self.model = embedding(model_name, cache_dir, threads, **kwargs)
|
||||
return
|
||||
break
|
||||
|
||||
raise ValueError(
|
||||
f"Model {model_name} is not supported in TextEmbedding."
|
||||
"Please check the supported models using `TextEmbedding.list_supported_models()`"
|
||||
)
|
||||
if self.model is None:
|
||||
raise ValueError(
|
||||
f"Model {model_name} is not supported in TextEmbedding."
|
||||
"Please check the supported models using `TextEmbedding.list_supported_models()`"
|
||||
)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
parallel: int = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
"""
|
||||
|
||||
@@ -19,7 +19,7 @@ class TextEmbeddingBase(ModelManagement):
|
||||
self,
|
||||
documents: Union[str, Iterable[str]],
|
||||
batch_size: int = 256,
|
||||
parallel: Optional[int] = None,
|
||||
parallel: int = None,
|
||||
**kwargs,
|
||||
) -> Iterable[np.ndarray]:
|
||||
raise NotImplementedError()
|
||||
|
||||
Generated
+389
-454
File diff suppressed because it is too large
Load Diff
+6
-7
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "fastembed"
|
||||
version = "0.2.2"
|
||||
version = "0.2.1"
|
||||
description = "Fast, light, accurate library built for retrieval embedding generation"
|
||||
authors = ["NirantK <nirant.bits@gmail.com>"]
|
||||
license = "Apache License"
|
||||
@@ -26,15 +26,14 @@ numpy = [
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = "^7.4.2"
|
||||
ruff = "^0.2.2"
|
||||
ruff = "^0.1.13"
|
||||
notebook = ">=7.0.2"
|
||||
mkdocs-material = "^9.5.10"
|
||||
mkdocstrings = "^0.24.0"
|
||||
pillow = "^10.2.0"
|
||||
mkdocs-material = "^9.1.21"
|
||||
mkdocstrings = "^0.22.0"
|
||||
pillow = "^10.0.0"
|
||||
cairosvg = "^2.7.1"
|
||||
mknotebooks = "^0.8.0"
|
||||
pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
|
||||
click = "^8.1.7"
|
||||
pre-commit = { version = "^3.6.0", python = ">=3.9,<3.12" }
|
||||
|
||||
|
||||
[build-system]
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from fastembed.embedding import DefaultEmbedding, JinaEmbedding
|
||||
|
||||
CANONICAL_VECTOR_VALUES = {
|
||||
"BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
|
||||
"BAAI/bge-small-en-v1.5": np.array([0.01522374, -0.02271799, 0.00860278, -0.07424029, 0.00386434]),
|
||||
"BAAI/bge-small-zh-v1.5": np.array([-0.01023294, 0.07634465, 0.0691722, -0.04458365, -0.03160762]),
|
||||
"BAAI/bge-base-en": np.array([0.0115, 0.0372, 0.0295, 0.0121, 0.0346]),
|
||||
"BAAI/bge-base-en-v1.5": np.array([0.01129394, 0.05493144, 0.02615099, 0.00328772, 0.02996045]),
|
||||
"BAAI/bge-large-en-v1.5": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
|
||||
"sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
|
||||
"intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
|
||||
"xenova/multilingual-e5-large": np.array([0.00975464, 0.00446568, 0.00655449, -0.0354155, 0.00702112]),
|
||||
"xenova/paraphrase-multilingual-mpnet-base-v2": np.array(
|
||||
[-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
|
||||
),
|
||||
"jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
|
||||
"jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("embedding_class", [DefaultEmbedding, JinaEmbedding])
|
||||
def test_embedding(embedding_class):
|
||||
is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
|
||||
|
||||
for model_desc in embedding_class.list_supported_models():
|
||||
if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
|
||||
continue
|
||||
|
||||
if model_desc["model"] not in CANONICAL_VECTOR_VALUES:
|
||||
continue
|
||||
|
||||
dim = model_desc["dim"]
|
||||
model = embedding_class(model_name=model_desc["model"])
|
||||
|
||||
docs = ["hello world", "flag embedding"]
|
||||
embeddings = list(model.embed(docs))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
assert embeddings.shape == (2, dim)
|
||||
|
||||
canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
|
||||
assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_dims,embedding_class", [(384, DefaultEmbedding), (768, JinaEmbedding)])
|
||||
def test_batch_embedding(n_dims, embedding_class):
|
||||
model = embedding_class()
|
||||
|
||||
docs = ["hello world", "flag embedding"] * 100
|
||||
embeddings = list(model.embed(docs, batch_size=10))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
assert embeddings.shape == (200, n_dims)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_dims,embedding_class", [(384, DefaultEmbedding), (768, JinaEmbedding)])
|
||||
def test_parallel_processing(n_dims, embedding_class):
|
||||
model = embedding_class()
|
||||
|
||||
docs = ["hello world", "flag embedding"] * 100
|
||||
embeddings = list(model.embed(docs, batch_size=10, parallel=2))
|
||||
embeddings = np.stack(embeddings, axis=0)
|
||||
|
||||
embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
|
||||
embeddings_2 = np.stack(embeddings_2, axis=0)
|
||||
|
||||
embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
|
||||
embeddings_3 = np.stack(embeddings_3, axis=0)
|
||||
|
||||
assert embeddings.shape == (200, n_dims)
|
||||
assert np.allclose(embeddings, embeddings_2, atol=1e-3)
|
||||
assert np.allclose(embeddings, embeddings_3, atol=1e-3)
|
||||
@@ -14,18 +14,12 @@ CANONICAL_VECTOR_VALUES = {
|
||||
"BAAI/bge-large-en-v1.5": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
|
||||
"BAAI/bge-large-en-v1.5-quantized": np.array([0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825]),
|
||||
"sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]),
|
||||
"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array([0.0094, 0.0184, 0.0328, 0.0072, -0.0351]),
|
||||
"intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
|
||||
"sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
|
||||
[-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
|
||||
),
|
||||
"jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]),
|
||||
"jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]),
|
||||
"nomic-ai/nomic-embed-text-v1": np.array([0.0061, 0.0103, -0.0296, -0.0242, -0.0170]),
|
||||
"nomic-ai/nomic-embed-text-v1.5": np.array(
|
||||
[-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
|
||||
),
|
||||
"thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user