Toggle Menu
Prev Up Next

# Source code for langchain_community.embeddings.itrex

    import importlib.util
    import os
    from typing import Any, Dict, List, Optional

    from langchain_core.embeddings import Embeddings
    from langchain_core.pydantic_v1 import BaseModel, Extra

    class QuantizedBgeEmbeddings(BaseModel, Embeddings):
        """Leverage Itrex runtime to unlock the performance of compressed NLP models.

        Please ensure that you have installed intel-extension-for-transformers.

        Input:
            model_name: str = Model name.
            max_seq_len: int = The maximum sequence length for tokenization. (default 512)
            pooling_strategy: str =
                "mean" or "cls", pooling strategy for the final layer. (default "mean")
            query_instruction: Optional[str] =
                An instruction to add to the query before embedding. (default None)
            document_instruction: Optional[str] =
                An instruction to add to each document before embedding. (default None)
            padding: Optional[bool] =
                Whether to add padding during tokenization or not. (default True)
            model_kwargs: Optional[Dict] =
                Parameters to add to the model during initialization. (default {})
            encode_kwargs: Optional[Dict] =
                Parameters to add during the embedding forward pass. (default {})
            onnx_file_name: Optional[str] =
                File name of onnx optimized model which is exported by itrex.
                (default "int8-model.onnx")

        Example:
            .. code-block:: python

                from langchain_community.embeddings import QuantizedBgeEmbeddings

                model_name = "Intel/bge-small-en-v1.5-sts-int8-static-inc"
                encode_kwargs = {'normalize_embeddings': True}
                hf = QuantizedBgeEmbeddings(
                    model_name,
                    encode_kwargs=encode_kwargs,
                    query_instruction="Represent this sentence for searching relevant passages: "
                )
        """  # noqa: E501

        def __init__(
            self,
            model_name: str,
            *,
            max_seq_len: int = 512,
            pooling_strategy: str = "mean",  # "mean" or "cls"
            query_instruction: Optional[str] = None,
            document_instruction: Optional[str] = None,
            padding: bool = True,
            model_kwargs: Optional[Dict] = None,
            encode_kwargs: Optional[Dict] = None,
            onnx_file_name: Optional[str] = "int8-model.onnx",
            **kwargs: Any,
        ) -> None:
            super().__init__(**kwargs)

            # check sentence_transformers python package
            if importlib.util.find_spec("intel_extension_for_transformers") is None:
                raise ImportError(
                    "Could not import intel_extension_for_transformers python package. "
                    "Please install it with "
                    "`pip install -U intel-extension-for-transformers`."
                )

            # check torch python package
            if importlib.util.find_spec("torch") is None:
                raise ImportError(
                    "Could not import torch python package. "
                    "Please install it with `pip install -U torch`."
                )

            # check onnx python package
            if importlib.util.find_spec("onnx") is None:
                raise ImportError(
                    "Could not import onnx python package. "
                    "Please install it with `pip install -U onnx`."
                )

            self.model_name_or_path = model_name
            self.max_seq_len = max_seq_len
            self.pooling = pooling_strategy
            self.padding = padding
            self.encode_kwargs = encode_kwargs or {}
            self.model_kwargs = model_kwargs or {}

            self.normalize = self.encode_kwargs.get("normalize_embeddings", False)
            self.batch_size = self.encode_kwargs.get("batch_size", 32)

            self.query_instruction = query_instruction
            self.document_instruction = document_instruction
            self.onnx_file_name = onnx_file_name

            self.load_model()

        def load_model(self) -> None:
            from huggingface_hub import hf_hub_download
            from intel_extension_for_transformers.transformers import AutoModel
            from transformers import AutoConfig, AutoTokenizer

            self.hidden_size = AutoConfig.from_pretrained(
                self.model_name_or_path
            ).hidden_size
            self.transformer_tokenizer = AutoTokenizer.from_pretrained(
                self.model_name_or_path,
            )
            onnx_model_path = os.path.join(self.model_name_or_path, self.onnx_file_name)  # type: ignore[arg-type]
            if not os.path.exists(onnx_model_path):
                onnx_model_path = hf_hub_download(
                    self.model_name_or_path, filename=self.onnx_file_name
                )
            self.transformer_model = AutoModel.from_pretrained(
                onnx_model_path, use_embedding_runtime=True
            )

        class Config:
            """Configuration for this pydantic object."""

            extra = Extra.allow

        def _embed(self, inputs: Any) -> Any:
            import torch

            engine_input = [value for value in inputs.values()]
            outputs = self.transformer_model.generate(engine_input)
            if "last_hidden_state:0" in outputs:
                last_hidden_state = outputs["last_hidden_state:0"]
            else:
                last_hidden_state = [out for out in outputs.values()][0]
            last_hidden_state = torch.tensor(last_hidden_state).reshape(
                inputs["input_ids"].shape[0], inputs["input_ids"].shape[1], self.hidden_size
            )
            if self.pooling == "mean":
                emb = self._mean_pooling(last_hidden_state, inputs["attention_mask"])
            elif self.pooling == "cls":
                emb = self._cls_pooling(last_hidden_state)
            else:
                raise ValueError("pooling method no supported")

            if self.normalize:
                emb = torch.nn.functional.normalize(emb, p=2, dim=1)
            return emb

        @staticmethod
        def _cls_pooling(last_hidden_state: Any) -> Any:
            return last_hidden_state[:, 0]

        @staticmethod
        def _mean_pooling(last_hidden_state: Any, attention_mask: Any) -> Any:
            try:
                import torch
            except ImportError as e:
                raise ImportError(
                    "Unable to import torch, please install with `pip install -U torch`."
                ) from e
            input_mask_expanded = (
                attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
            )
            sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)
            sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
            return sum_embeddings / sum_mask

        def _embed_text(self, texts: List[str]) -> List[List[float]]:
            inputs = self.transformer_tokenizer(
                texts,
                max_length=self.max_seq_len,
                truncation=True,
                padding=self.padding,
                return_tensors="pt",
            )
            return self._embed(inputs).tolist()

        def embed_documents(self, texts: List[str]) -> List[List[float]]:
            """Embed a list of text documents using the Optimized Embedder model.

            Input:
                texts: List[str] = List of text documents to embed.
            Output:
                List[List[float]] = The embeddings of each text document.
            """
            try:
                import pandas as pd
            except ImportError as e:
                raise ImportError(
                    "Unable to import pandas, please install with `pip install -U pandas`."
                ) from e
            docs = [
                self.document_instruction + d if self.document_instruction else d
                for d in texts
            ]

            # group into batches
            text_list_df = pd.DataFrame(docs, columns=["texts"]).reset_index()

            # assign each example with its batch
            text_list_df["batch_index"] = text_list_df["index"] // self.batch_size

            # create groups
            batches = list(text_list_df.groupby(["batch_index"])["texts"].apply(list))

            vectors = []
            for batch in batches:
                vectors += self._embed_text(batch)
            return vectors

        def embed_query(self, text: str) -> List[float]:
            if self.query_instruction:
                text = self.query_instruction + text
            return self._embed_text([text])[0]