codekingpro/portable-devtools

1from pathlib import Path2from typing import Any, Dict, List3 4from langchain_core.embeddings import Embeddings5from pydantic import BaseModel, ConfigDict, Field6 7DEFAULT_QUERY_INSTRUCTION = (8    "Represent the question for retrieving supporting documents: "9)10DEFAULT_QUERY_BGE_INSTRUCTION_EN = (11    "Represent this question for searching relevant passages: "12)13DEFAULT_QUERY_BGE_INSTRUCTION_ZH = "为这个句子生成表示以用于检索相关文章："14 15 16class OpenVINOEmbeddings(BaseModel, Embeddings):17    """OpenVINO embedding models.18 19    Example:20        .. code-block:: python21 22            from langchain_community.embeddings import OpenVINOEmbeddings23 24            model_name = "sentence-transformers/all-mpnet-base-v2"25            model_kwargs = {'device': 'CPU'}26            encode_kwargs = {'normalize_embeddings': True}27            ov = OpenVINOEmbeddings(28                model_name_or_path=model_name,29                model_kwargs=model_kwargs,30                encode_kwargs=encode_kwargs31            )32    """33 34    ov_model: Any = None35    """OpenVINO model object."""36    tokenizer: Any = None37    """Tokenizer for embedding model."""38    model_name_or_path: str39    """HuggingFace model id."""40    model_kwargs: Dict[str, Any] = Field(default_factory=dict)41    """Keyword arguments to pass to the model."""42    encode_kwargs: Dict[str, Any] = Field(default_factory=dict)43    """Keyword arguments to pass when calling the `encode` method of the model."""44    show_progress: bool = False45    """Whether to show a progress bar."""46 47    def __init__(self, **kwargs: Any):48        """Initialize the sentence_transformer."""49        super().__init__(**kwargs)50 51        try:52            from optimum.intel.openvino import OVModelForFeatureExtraction53        except ImportError as e:54            raise ImportError(55                "Could not import optimum-intel python package. "56                "Please install it with: "57                "pip install -U 'optimum[openvino,nncf]'"58            ) from e59 60        try:61            from huggingface_hub import HfApi62        except ImportError as e:63            raise ImportError(64                "Could not import huggingface_hub python package. "65                "Please install it with: "66                "`pip install -U huggingface_hub`."67            ) from e68 69        def require_model_export(70            model_id: str, revision: Any = None, subfolder: Any = None71        ) -> bool:72            model_dir = Path(model_id)73            if subfolder is not None:74                model_dir = model_dir / subfolder75            if model_dir.is_dir():76                return (77                    not (model_dir / "openvino_model.xml").exists()78                    or not (model_dir / "openvino_model.bin").exists()79                )80            hf_api = HfApi()81            try:82                model_info = hf_api.model_info(model_id, revision=revision or "main")83                normalized_subfolder = (84                    None if subfolder is None else Path(subfolder).as_posix()85                )86                model_files = [87                    file.rfilename88                    for file in model_info.siblings89                    if normalized_subfolder is None90                    or file.rfilename.startswith(normalized_subfolder)91                ]92                ov_model_path = (93                    "openvino_model.xml"94                    if subfolder is None95                    else f"{normalized_subfolder}/openvino_model.xml"96                )97                return (98                    ov_model_path not in model_files99                    or ov_model_path.replace(".xml", ".bin") not in model_files100                )101            except Exception:102                return True103 104        if require_model_export(self.model_name_or_path):105            # use remote model106            self.ov_model = OVModelForFeatureExtraction.from_pretrained(107                self.model_name_or_path, export=True, **self.model_kwargs108            )109        else:110            # use local model111            self.ov_model = OVModelForFeatureExtraction.from_pretrained(112                self.model_name_or_path, **self.model_kwargs113            )114 115        try:116            from transformers import AutoTokenizer117        except ImportError as e:118            raise ImportError(119                "Unable to import transformers, please install with "120                "`pip install -U transformers`."121            ) from e122        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)123 124    def _text_length(self, text: Any) -> int:125        """126        Help function to get the length for the input text. Text can be either127        a list of ints (which means a single text as input), or a tuple of list of ints128        (representing several text inputs to the model).129        """130 131        if isinstance(text, dict):  # {key: value} case132            return len(next(iter(text.values())))133        elif not hasattr(text, "__len__"):  # Object has no len() method134            return 1135        # Empty string or list of ints136        elif len(text) == 0 or isinstance(text[0], int):137            return len(text)138        else:139            # Sum of length of individual strings140            return sum([len(t) for t in text])141 142    def encode(143        self,144        sentences: Any,145        batch_size: int = 4,146        show_progress_bar: bool = False,147        convert_to_numpy: bool = True,148        convert_to_tensor: bool = False,149        mean_pooling: bool = False,150        normalize_embeddings: bool = True,151    ) -> Any:152        """153        Computes sentence embeddings.154 155        :param sentences: the sentences to embed.156        :param batch_size: the batch size used for the computation.157        :param show_progress_bar: Whether to output a progress bar.158        :param convert_to_numpy: Whether the output should be a list of numpy vectors.159        :param convert_to_tensor: Whether the output should be one large tensor.160        :param mean_pooling: Whether to pool returned vectors.161        :param normalize_embeddings: Whether to normalize returned vectors.162 163        :return: By default, a 2d numpy array with shape [num_inputs, output_dimension].164        """165        try:166            import numpy as np167        except ImportError as e:168            raise ImportError(169                "Unable to import numpy, please install with `pip install -U numpy`."170            ) from e171        try:172            from tqdm import trange173        except ImportError as e:174            raise ImportError(175                "Unable to import tqdm, please install with `pip install -U tqdm`."176            ) from e177        try:178            import torch179        except ImportError as e:180            raise ImportError(181                "Unable to import torch, please install with `pip install -U torch`."182            ) from e183 184        def run_mean_pooling(model_output: Any, attention_mask: Any) -> Any:185            token_embeddings = model_output[186                0187            ]  # First element of model_output contains all token embeddings188            input_mask_expanded = (189                attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()190            )191            return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(192                input_mask_expanded.sum(1), min=1e-9193            )194 195        if convert_to_tensor:196            convert_to_numpy = False197 198        input_was_string = False199        if isinstance(sentences, str) or not hasattr(200            sentences, "__len__"201        ):  # Cast an individual sentence to a list with length 1202            sentences = [sentences]203            input_was_string = True204 205        all_embeddings: Any = []206        length_sorted_idx = np.argsort([-self._text_length(sen) for sen in sentences])207        sentences_sorted = [sentences[idx] for idx in length_sorted_idx]208 209        for start_index in trange(210            0, len(sentences), batch_size, desc="Batches", disable=not show_progress_bar211        ):212            sentences_batch = sentences_sorted[start_index : start_index + batch_size]213 214            length = self.ov_model.request.inputs[0].get_partial_shape()[1]215            if length.is_dynamic:216                features = self.tokenizer(217                    sentences_batch, padding=True, truncation=True, return_tensors="pt"218                )219            else:220                features = self.tokenizer(221                    sentences_batch,222                    padding="max_length",223                    max_length=length.get_length(),224                    truncation=True,225                    return_tensors="pt",226                )227 228            out_features = self.ov_model(**features)229            if mean_pooling:230                embeddings = run_mean_pooling(out_features, features["attention_mask"])231            else:232                embeddings = out_features[0][:, 0]233            if normalize_embeddings:234                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)235 236            # fixes for #522 and #487 to avoid oom problems on gpu with large datasets237            if convert_to_numpy:238                embeddings = embeddings.cpu()239 240            all_embeddings.extend(embeddings)241 242        all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]243 244        if convert_to_tensor:245            if len(all_embeddings):246                all_embeddings = torch.stack(all_embeddings)247            else:248                all_embeddings = torch.Tensor()249        elif convert_to_numpy:250            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])251 252        if input_was_string:253            all_embeddings = all_embeddings[0]254 255        return all_embeddings256 257    model_config = ConfigDict(extra="forbid", protected_namespaces=())258 259    def embed_documents(self, texts: List[str]) -> List[List[float]]:260        """Compute doc embeddings using a HuggingFace transformer model.261 262        Args:263            texts: The list of texts to embed.264 265        Returns:266            List of embeddings, one for each text.267        """268 269        texts = list(map(lambda x: x.replace("\n", " "), texts))270        embeddings = self.encode(271            texts, show_progress_bar=self.show_progress, **self.encode_kwargs272        )273 274        return embeddings.tolist()275 276    def embed_query(self, text: str) -> List[float]:277        """Compute query embeddings using a HuggingFace transformer model.278 279        Args:280            text: The text to embed.281 282        Returns:283            Embeddings for the text.284        """285        return self.embed_documents([text])[0]286 287    def save_model(288        self,289        model_path: str,290    ) -> bool:291        self.ov_model.half()292        self.ov_model.save_pretrained(model_path)293        self.tokenizer.save_pretrained(model_path)294        return True295 296 297class OpenVINOBgeEmbeddings(OpenVINOEmbeddings):298    """OpenVNO BGE embedding models.299 300    Bge Example:301        .. code-block:: python302 303            from langchain_community.embeddings import OpenVINOBgeEmbeddings304 305            model_name = "BAAI/bge-large-en-v1.5"306            model_kwargs = {'device': 'CPU'}307            encode_kwargs = {'normalize_embeddings': True}308            ov = OpenVINOBgeEmbeddings(309                model_name_or_path=model_name,310                model_kwargs=model_kwargs,311                encode_kwargs=encode_kwargs312            )313    """314 315    query_instruction: str = DEFAULT_QUERY_BGE_INSTRUCTION_EN316    """Instruction to use for embedding query."""317    embed_instruction: str = ""318    """Instruction to use for embedding document."""319 320    def __init__(self, **kwargs: Any):321        """Initialize the sentence_transformer."""322        super().__init__(**kwargs)323 324        if "-zh" in self.model_name_or_path:325            self.query_instruction = DEFAULT_QUERY_BGE_INSTRUCTION_ZH326 327    def embed_documents(self, texts: List[str]) -> List[List[float]]:328        """Compute doc embeddings using a HuggingFace transformer model.329 330        Args:331            texts: The list of texts to embed.332 333        Returns:334            List of embeddings, one for each text.335        """336        texts = [self.embed_instruction + t.replace("\n", " ") for t in texts]337        embeddings = self.encode(texts, **self.encode_kwargs)338        return embeddings.tolist()339 340    def embed_query(self, text: str) -> List[float]:341        """Compute query embeddings using a HuggingFace transformer model.342 343        Args:344            text: The text to embed.345 346        Returns:347            Embeddings for the text.348        """349        text = text.replace("\n", " ")350        embedding = self.encode(self.query_instruction + text, **self.encode_kwargs)351        return embedding.tolist()352