Add dynamic quantization from @Xenova

* check for existing files

* (chore: model_management.py):  Add comment

---------

Co-authored-by: Nirant <NirantK@users.noreply.github.com>

docs/examples/Supported_Models.ipynb

@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -114,10 +114,26 @@
        "      <td>768</td>\n",
        "      <td>8192 context length english model</td>\n",
        "      <td>0.54</td>\n",
-       "      <td>{'hf': 'xenova/nomic-embed-text-v1'}</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",
@@ -125,7 +141,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>10</th>\n",
+       "      <th>12</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",
@@ -133,7 +149,15 @@
        "      <td>{'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>11</th>\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",
        "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
        "      <td>768</td>\n",
        "      <td>English embedding model supporting 8192 sequence length</td>\n",
@@ -141,7 +165,7 @@
        "      <td>{'hf': 'xenova/jina-embeddings-v2-base-en'}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
-       "      <th>12</th>\n",
+       "      <th>15</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",
@@ -163,10 +187,13 @@
        "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                                intfloat/multilingual-e5-large  1024   \n",
-       "10  sentence-transformers/paraphrase-multilingual-mpnet-base-v2   768   \n",
-       "11                            jinaai/jina-embeddings-v2-base-en   768   \n",
-       "12                           jinaai/jina-embeddings-v2-small-en   512   \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",
        "\n",
        "                                                                           description  \\\n",
        "0                                                                   Base English model   \n",
@@ -178,10 +205,13 @@
        "6                                                   Fast and recommended Chinese model   \n",
        "7                                             Sentence Transformer model, MiniLM-L6-v2   \n",
        "8                                                    8192 context length english model   \n",
-       "9   Multilingual model, e5-large. Recommend using this model for non-English languages   \n",
-       "10            Sentence-transformers model for tasks like clustering or semantic search   \n",
-       "11                             English embedding model supporting 8192 sequence length   \n",
-       "12                             English embedding model supporting 8192 sequence length   \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",
        "\n",
        "    size_in_GB  \\\n",
        "0         0.50   \n",
@@ -193,10 +223,13 @@
        "6         0.10   \n",
        "7         0.09   \n",
        "8         0.54   \n",
-       "9         2.24   \n",
-       "10        1.11   \n",
-       "11        0.55   \n",
-       "12        0.13   \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",
        "\n",
        "                                                                                                                                           sources  \n",
        "0                                                               {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz'}  \n",
@@ -207,14 +240,17 @@
        "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': 'xenova/nomic-embed-text-v1'}  \n",
-       "9          {'url': 'https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz', 'hf': 'qdrant/multilingual-e5-large-onnx'}  \n",
-       "10                                                                                          {'hf': 'xenova/paraphrase-multilingual-mpnet-base-v2'}  \n",
-       "11                                                                                                     {'hf': 'xenova/jina-embeddings-v2-base-en'}  \n",
-       "12                                                                                                    {'hf': 'xenova/jina-embeddings-v2-small-en'}  "
+       "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'}  "
       ]
      },
-     "execution_count": 4,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -244,7 +280,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.10.13"
   },
   "orig_nbformat": 4
  },

experiments/dynamic_quantization_export.py

@@ -0,0 +1,473 @@
+"""
+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()

experiments/try_quantized_model.py

@@ -0,0 +1,56 @@
+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()

fastembed/common/model_management.py

@@ -131,7 +131,8 @@ class ModelManagement:
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
 
-        if model_dir.exists():
+        # 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:
             return model_dir
 
         if model_tmp_dir.exists():

fastembed/models.json

@@ -1,133 +0,0 @@
-[
-    {
-        "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": []
-    }
-]

fastembed/text/onnx_embedding.py

@@ -98,6 +98,15 @@ 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,
@@ -107,6 +116,24 @@ supported_onnx_models = [
             "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,

poetry.lock

@@ -3430,4 +3430,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.0,<3.13"
-content-hash = "04730c4b2a45b7a1ffc5faf449a2faedbed38097aae9cfd0dfa6ce533197db80"
+content-hash = "8bc18a20773e3e2530564b9e0ba8a31c6f3dd7ab710c87acfe446db5e785ca3c"

pyproject.toml

@@ -34,6 +34,7 @@ pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
+click = "^8.1.7"
 
 
 [build-system]

tests/test_onnx_embeddings.py

@@ -1,77 +0,0 @@
-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]),
-    "nomic-ai/nomic-embed-text-v1": np.array([0.0061, 0.0103, -0.0296, -0.0242, -0.0170]),
-}
-
-
-@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)

tests/test_text_onnx_embeddings.py

@@ -14,6 +14,7 @@ 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]
@@ -21,6 +22,10 @@ CANONICAL_VECTOR_VALUES = {
     "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]),
 }