From 0f85dea8c89f00a72ef52043a6d2abad1cd0d1ae Mon Sep 17 00:00:00 2001 From: Roman Solomatin Date: Thu, 7 Nov 2024 03:35:39 +0500 Subject: [PATCH] langchain-huggingface: use separate kwargs for queries and docs (#27857) Now `encode_kwargs` used for both for documents and queries and this leads to wrong embeddings. E. g.: ```python model_kwargs = {"device": "cuda", "trust_remote_code": True} encode_kwargs = {"normalize_embeddings": False, "prompt_name": "s2p_query"} model = HuggingFaceEmbeddings( model_name="dunzhang/stella_en_400M_v5", model_kwargs=model_kwargs, encode_kwargs=encode_kwargs, ) query_embedding = np.array( model.embed_query("What are some ways to reduce stress?",) ) document_embedding = np.array( model.embed_documents( [ "There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.", "Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.", ] ) ) print(model._client.similarity(query_embedding, document_embedding)) # output: tensor([[0.8421, 0.3317]], dtype=torch.float64) ``` But from the [model card](https://huggingface.co/dunzhang/stella_en_400M_v5#sentence-transformers) expexted like this: ```python model_kwargs = {"device": "cuda", "trust_remote_code": True} encode_kwargs = {"normalize_embeddings": False} query_encode_kwargs = {"normalize_embeddings": False, "prompt_name": "s2p_query"} model = HuggingFaceEmbeddings( model_name="dunzhang/stella_en_400M_v5", model_kwargs=model_kwargs, encode_kwargs=encode_kwargs, query_encode_kwargs=query_encode_kwargs, ) query_embedding = np.array( model.embed_query("What are some ways to reduce stress?", ) ) document_embedding = np.array( model.embed_documents( [ "There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.", "Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.", ] ) ) print(model._client.similarity(query_embedding, document_embedding)) # tensor([[0.8398, 0.2990]], dtype=torch.float64) ``` --- .../embeddings/huggingface.py | 41 +++++++++++++++---- 1 file changed, 34 insertions(+), 7 deletions(-) diff --git a/libs/partners/huggingface/langchain_huggingface/embeddings/huggingface.py b/libs/partners/huggingface/langchain_huggingface/embeddings/huggingface.py index 180a9ed3b5e..2bbc551f4e0 100644 --- a/libs/partners/huggingface/langchain_huggingface/embeddings/huggingface.py +++ b/libs/partners/huggingface/langchain_huggingface/embeddings/huggingface.py @@ -36,9 +36,14 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings): `prompts`, `default_prompt_name`, `revision`, `trust_remote_code`, or `token`. See also the Sentence Transformer documentation: https://sbert.net/docs/package_reference/SentenceTransformer.html#sentence_transformers.SentenceTransformer""" encode_kwargs: Dict[str, Any] = Field(default_factory=dict) - """Keyword arguments to pass when calling the `encode` method of the Sentence - Transformer model, such as `prompt_name`, `prompt`, `batch_size`, `precision`, - `normalize_embeddings`, and more. + """Keyword arguments to pass when calling the `encode` method for the documents of + the Sentence Transformer model, such as `prompt_name`, `prompt`, `batch_size`, + `precision`, `normalize_embeddings`, and more. + See also the Sentence Transformer documentation: https://sbert.net/docs/package_reference/SentenceTransformer.html#sentence_transformers.SentenceTransformer.encode""" + query_encode_kwargs: Dict[str, Any] = Field(default_factory=dict) + """Keyword arguments to pass when calling the `encode` method for the query of + the Sentence Transformer model, such as `prompt_name`, `prompt`, `batch_size`, + `precision`, `normalize_embeddings`, and more. See also the Sentence Transformer documentation: https://sbert.net/docs/package_reference/SentenceTransformer.html#sentence_transformers.SentenceTransformer.encode""" multi_process: bool = False """Run encode() on multiple GPUs.""" @@ -65,11 +70,17 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings): protected_namespaces=(), ) - def embed_documents(self, texts: List[str]) -> List[List[float]]: - """Compute doc embeddings using a HuggingFace transformer model. + def _embed( + self, texts: list[str], encode_kwargs: Dict[str, Any] + ) -> List[List[float]]: + """ + Embed a text using the HuggingFace transformer model. Args: texts: The list of texts to embed. + encode_kwargs: Keyword arguments to pass when calling the + `encode` method for the documents of the SentenceTransformer + encode method. Returns: List of embeddings, one for each text. @@ -85,7 +96,7 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings): embeddings = self._client.encode( texts, show_progress_bar=self.show_progress, - **self.encode_kwargs, # type: ignore + **encode_kwargs, # type: ignore ) if isinstance(embeddings, list): @@ -96,6 +107,17 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings): return embeddings.tolist() + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Compute doc embeddings using a HuggingFace transformer model. + + Args: + texts: The list of texts to embed. + + Returns: + List of embeddings, one for each text. + """ + return self._embed(texts, self.encode_kwargs) + def embed_query(self, text: str) -> List[float]: """Compute query embeddings using a HuggingFace transformer model. @@ -105,4 +127,9 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings): Returns: Embeddings for the text. """ - return self.embed_documents([text])[0] + embed_kwargs = ( + self.query_encode_kwargs + if len(self.query_encode_kwargs) > 0 + else self.encode_kwargs + ) + return self._embed([text], embed_kwargs)[0]