feat!: Add support for Optimum optimizers and quantizers (#496)

* feat!: Add support for Optimum optimizers and quantizers

* Fix docstring typo

* Apply suggestions from code review

Co-authored-by: Silvano Cerza <[email protected]>

---------

Co-authored-by: Silvano Cerza <[email protected]>

.github/workflows/optimum.yml

@@ -52,9 +52,9 @@ jobs:
         if: matrix.python-version == '3.9' && runner.os == 'Linux'
         run: hatch run lint:all
 
-      - name: Generate docs
-        if: matrix.python-version == '3.9' && runner.os == 'Linux'
-        run: hatch run docs
+      # - name: Generate docs
+      #   if: matrix.python-version == '3.9' && runner.os == 'Linux'
+      #   run: hatch run docs
 
       - name: Run tests
         run: hatch run cov

integrations/optimum/pydoc/config.yml

@@ -6,6 +6,8 @@ loaders:
         "haystack_integrations.components.embedders.optimum.optimum_document_embedder",
         "haystack_integrations.components.embedders.optimum.optimum_text_embedder",
         "haystack_integrations.components.embedders.optimum.pooling",
+        "haystack_integrations.components.embedders.optimum.optimization",
+        "haystack_integrations.components.embedders.optimum.quantization",
       ]
     ignore_when_discovered: ["__init__"]
 processors:

integrations/optimum/src/haystack_integrations/components/embedders/optimum/__init__.py

@@ -2,8 +2,18 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from .optimization import OptimumEmbedderOptimizationConfig, OptimumEmbedderOptimizationMode
 from .optimum_document_embedder import OptimumDocumentEmbedder
 from .optimum_text_embedder import OptimumTextEmbedder
 from .pooling import OptimumEmbedderPooling
+from .quantization import OptimumEmbedderQuantizationConfig, OptimumEmbedderQuantizationMode
 
-__all__ = ["OptimumDocumentEmbedder", "OptimumEmbedderPooling", "OptimumTextEmbedder"]
+__all__ = [
+    "OptimumDocumentEmbedder",
+    "OptimumEmbedderOptimizationMode",
+    "OptimumEmbedderOptimizationConfig",
+    "OptimumEmbedderPooling",
+    "OptimumEmbedderQuantizationMode",
+    "OptimumEmbedderQuantizationConfig",
+    "OptimumTextEmbedder",
+]

integrations/optimum/src/haystack_integrations/components/embedders/optimum/_backend.py

@@ -1,7 +1,8 @@
 import copy
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -11,10 +12,17 @@
 from sentence_transformers.models import Pooling as SentenceTransformerPoolingLayer
 from tqdm import tqdm
 from transformers import AutoTokenizer
+from transformers.modeling_outputs import BaseModelOutput
 
-from optimum.onnxruntime import ORTModelForFeatureExtraction
+from optimum.onnxruntime import (
+    ORTModelForFeatureExtraction,
+    ORTOptimizer,
+    ORTQuantizer,
+)
 
+from .optimization import OptimumEmbedderOptimizationConfig
 from .pooling import OptimumEmbedderPooling
+from .quantization import OptimumEmbedderQuantizationConfig
 
 
 @dataclass
@@ -29,23 +37,41 @@ class _EmbedderParams:
     progress_bar: bool
     pooling_mode: Optional[Union[str, OptimumEmbedderPooling]]
     model_kwargs: Optional[Dict[str, Any]]
+    working_dir: Optional[str]
+    optimizer_settings: Optional[OptimumEmbedderOptimizationConfig]
+    quantizer_settings: Optional[OptimumEmbedderQuantizationConfig]
 
     def serialize(self) -> Dict[str, Any]:
         out = {}
         for field in self.__dataclass_fields__.keys():
+            if field in [
+                "pooling_mode",
+                "token",
+                "optimizer_settings",
+                "quantizer_settings",
+            ]:
+                continue
             out[field] = copy.deepcopy(getattr(self, field))
 
         # Fixups.
         assert isinstance(self.pooling_mode, OptimumEmbedderPooling)
-        out["pooling_mode"] = self.pooling_mode.value
+        out["pooling_mode"] = str(self.pooling_mode)
         out["token"] = self.token.to_dict() if self.token else None
+        out["optimizer_settings"] = self.optimizer_settings.to_dict() if self.optimizer_settings else None
+        out["quantizer_settings"] = self.quantizer_settings.to_dict() if self.quantizer_settings else None
+
         out["model_kwargs"].pop("use_auth_token", None)
         serialize_hf_model_kwargs(out["model_kwargs"])
         return out
 
     @classmethod
     def deserialize_inplace(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         data["pooling_mode"] = OptimumEmbedderPooling.from_str(data["pooling_mode"])
+        if data["optimizer_settings"] is not None:
+            data["optimizer_settings"] = OptimumEmbedderOptimizationConfig.from_dict(data["optimizer_settings"])
+        if data["quantizer_settings"] is not None:
+            data["quantizer_settings"] = OptimumEmbedderQuantizationConfig.from_dict(data["quantizer_settings"])
+
         deserialize_secrets_inplace(data, keys=["token"])
         deserialize_hf_model_kwargs(data["model_kwargs"])
         return data
@@ -71,6 +97,11 @@ def __init__(self, params: _EmbedderParams):
 
         params.model_kwargs = params.model_kwargs or {}
 
+        if params.optimizer_settings or params.quantizer_settings:
+            if not params.working_dir:
+                msg = "Working directory is required for optimization and quantization"
+                raise ValueError(msg)
+
         # Check if the model_kwargs contain the parameters, otherwise, populate them with values from init parameters
         params.model_kwargs.setdefault("model_id", params.model)
         params.model_kwargs.setdefault("provider", params.onnx_execution_provider)
@@ -82,18 +113,48 @@ def __init__(self, params: _EmbedderParams):
         self.pooling_layer = None
 
     def warm_up(self):
-        self.model = ORTModelForFeatureExtraction.from_pretrained(**self.params.model_kwargs, export=True)
+        assert self.params.model_kwargs
+        model_kwargs = copy.deepcopy(self.params.model_kwargs)
+        model = ORTModelForFeatureExtraction.from_pretrained(**model_kwargs, export=True)
+
+        # Model ID will be passed explicitly if optimization/quantization is enabled.
+        model_kwargs.pop("model_id", None)
+
+        optimized_model = False
+        if self.params.optimizer_settings:
+            assert self.params.working_dir
+            optimizer = ORTOptimizer.from_pretrained(model)
+            save_dir = optimizer.optimize(
+                save_dir=self.params.working_dir, optimization_config=self.params.optimizer_settings.to_optimum_config()
+            )
+            model = ORTModelForFeatureExtraction.from_pretrained(model_id=save_dir, **model_kwargs)
+            optimized_model = True
+
+        if self.params.quantizer_settings:
+            assert self.params.working_dir
+
+            # We need to create a subfolder for models that were optimized before quantization
+            # since Optimum expects no more than one ONXX model in the working directory. There's
+            # a file name parameter, but the optimizer only returns the working directory.
+            working_dir = (
+                Path(self.params.working_dir) if not optimized_model else Path(self.params.working_dir) / "quantized"
+            )
+            quantizer = ORTQuantizer.from_pretrained(model)
+            save_dir = quantizer.quantize(
+                save_dir=working_dir, quantization_config=self.params.quantizer_settings.to_optimum_config()
+            )
+            model = ORTModelForFeatureExtraction.from_pretrained(model_id=save_dir, **model_kwargs)
+
+        self.model = model
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.params.model, token=self.params.token.resolve_value() if self.params.token else None
         )
 
         # We need the width of the embeddings to initialize the pooling layer
         # so we do a dummy forward pass with the model.
-        dummy_input = self.tokenizer(["dummy input"], padding=True, truncation=True, return_tensors="pt").to(
-            self.model.device
-        )
-        dummy_output = self.model(input_ids=dummy_input["input_ids"], attention_mask=dummy_input["attention_mask"])
-        width = dummy_output[0].size(dim=2)  # BaseModelOutput.last_hidden_state
+        width = self._tokenize_and_generate_outputs(["dummy input"])[1][0].size(
+            dim=2
+        )  # BaseModelOutput.last_hidden_state
 
         self.pooling_layer = SentenceTransformerPoolingLayer(
             width,
@@ -105,6 +166,17 @@ def warm_up(self):
             pooling_mode_lasttoken=self.params.pooling_mode == OptimumEmbedderPooling.LAST_TOKEN,
         )
 
+    def _tokenize_and_generate_outputs(self, texts: List[str]) -> Tuple[Dict[str, Any], BaseModelOutput]:
+        assert self.model is not None
+        assert self.tokenizer is not None
+
+        tokenizer_outputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(
+            self.model.device
+        )
+        model_inputs = {k: v for k, v in tokenizer_outputs.items() if k in self.model.inputs_names}
+        model_outputs = self.model(**model_inputs)
+        return tokenizer_outputs, model_outputs
+
     @property
     def parameters(self) -> _EmbedderParams:
         return self.params
@@ -140,11 +212,8 @@ def embed_texts(
             desc="Calculating embeddings",
         ):
             batch = sentences_sorted[i : i + self.params.batch_size]
-            encoded_input = self.tokenizer(batch, padding=True, truncation=True, return_tensors="pt").to(device)
-            model_output = self.model(
-                input_ids=encoded_input["input_ids"], attention_mask=encoded_input["attention_mask"]
-            )
-            sentence_embeddings = self.pool_embeddings(model_output[0], encoded_input["attention_mask"].to(device))
+            tokenizer_output, model_output = self._tokenize_and_generate_outputs(batch)
+            sentence_embeddings = self.pool_embeddings(model_output[0], tokenizer_output["attention_mask"].to(device))
             all_embeddings.append(sentence_embeddings)
 
         embeddings = torch.cat(all_embeddings, dim=0)

integrations/optimum/src/haystack_integrations/components/embedders/optimum/optimization.py

@@ -0,0 +1,105 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict
+
+from optimum.onnxruntime.configuration import AutoOptimizationConfig, OptimizationConfig
+
+
+class OptimumEmbedderOptimizationMode(Enum):
+    """
+    [ONXX Optimization Modes](https://huggingface.co/docs/optimum/onnxruntime/usage_guides/optimization.html)
+    support by the Optimum Embedders.
+    """
+
+    #: Basic general optimizations.
+    O1 = "o1"
+
+    #: Basic and extended general optimizations, transformers-specific fusions.
+    O2 = "o2"
+
+    #: Same as O2 with Gelu approximation.
+    O3 = "o3"
+
+    #: Same as O3 with mixed precision.
+    O4 = "o4"
+
+    def __str__(self):
+        return self.value
+
+    @classmethod
+    def from_str(cls, string: str) -> "OptimumEmbedderOptimizationMode":
+        """
+        Create an optimization mode from a string.
+
+        :param string:
+            String to convert.
+        :returns:
+            Optimization mode.
+        """
+        enum_map = {e.value: e for e in OptimumEmbedderOptimizationMode}
+        opt_mode = enum_map.get(string)
+        if opt_mode is None:
+            msg = f"Unknown optimization mode '{string}'. Supported modes are: {list(enum_map.keys())}"
+            raise ValueError(msg)
+        return opt_mode
+
+
+@dataclass(frozen=True)
+class OptimumEmbedderOptimizationConfig:
+    """
+    Configuration for Optimum Embedder Optimization.
+
+    :param mode:
+        Optimization mode.
+    :param for_gpu:
+        Whether to optimize for GPUs.
+    """
+
+    mode: OptimumEmbedderOptimizationMode
+    for_gpu: bool = True
+
+    def to_optimum_config(self) -> OptimizationConfig:
+        """
+        Convert the configuration to a Optimum configuration.
+
+        :returns:
+            Optimum configuration.
+        """
+        if self.mode == OptimumEmbedderOptimizationMode.O1:
+            return AutoOptimizationConfig.O1(for_gpu=self.for_gpu)
+        elif self.mode == OptimumEmbedderOptimizationMode.O2:
+            return AutoOptimizationConfig.O2(for_gpu=self.for_gpu)
+        elif self.mode == OptimumEmbedderOptimizationMode.O3:
+            return AutoOptimizationConfig.O3(for_gpu=self.for_gpu)
+        elif self.mode == OptimumEmbedderOptimizationMode.O4:
+            return AutoOptimizationConfig.O4(for_gpu=self.for_gpu)
+        else:
+            msg = f"Unknown optimization mode '{self.mode}'"
+            raise ValueError(msg)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the configuration to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
+        return {
+            "mode": str(self.mode),
+            "for_gpu": self.for_gpu,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "OptimumEmbedderOptimizationConfig":
+        """
+        Create an optimization configuration from a dictionary.
+
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Optimization configuration.
+        """
+        return OptimumEmbedderOptimizationConfig(
+            mode=OptimumEmbedderOptimizationMode.from_str(data["mode"]),
+            for_gpu=data["for_gpu"],
+        )