Files
langchain/libs/community/langchain_community/embeddings/tensorflow_hub.py
Eugene Yurtsev 844955d6e1 community[patch]: assign missed default (#26326)
Assigning missed defaults in various classes. Most clients were being
assigned during the `model_validator(mode="before")` step, so this
change should amount to a no-op in those cases.

---

This PR was autogenerated using gritql

```shell

grit apply 'class_definition(name=$C, $body, superclasses=$S) where {    
    $C <: ! "Config", // Does not work in this scope, but works after class_definition
    $body <: block($statements),
    $statements <: some bubble assignment(left=$x, right=$y, type=$t) as $A where {
        or {
            $y <: `Field($z)`,
            $x <: "model_config"
        }
    },
    // And has either Any or Optional fields without a default
    $statements <: some bubble assignment(left=$x, right=$y, type=$t) as $A where {
        $t <: or {
            r"Optional.*",
            r"Any",
            r"Union[None, .*]",
            r"Union[.*, None, .*]",
            r"Union[.*, None]",
        },
        $y <: ., // Match empty node        
        $t => `$t = None`,
    },    
}
' --language python .

```
2024-09-11 11:13:11 -04:00

75 lines
2.3 KiB
Python

from typing import Any, List
from langchain_core.embeddings import Embeddings
from pydantic import BaseModel, ConfigDict
DEFAULT_MODEL_URL = "https://tfhub.dev/google/universal-sentence-encoder-multilingual/3"
class TensorflowHubEmbeddings(BaseModel, Embeddings):
"""TensorflowHub embedding models.
To use, you should have the ``tensorflow_text`` python package installed.
Example:
.. code-block:: python
from langchain_community.embeddings import TensorflowHubEmbeddings
url = "https://tfhub.dev/google/universal-sentence-encoder-multilingual/3"
tf = TensorflowHubEmbeddings(model_url=url)
"""
embed: Any = None #: :meta private:
model_url: str = DEFAULT_MODEL_URL
"""Model name to use."""
def __init__(self, **kwargs: Any):
"""Initialize the tensorflow_hub and tensorflow_text."""
super().__init__(**kwargs)
try:
import tensorflow_hub
except ImportError:
raise ImportError(
"Could not import tensorflow-hub python package. "
"Please install it with `pip install tensorflow-hub``."
)
try:
import tensorflow_text # noqa
except ImportError:
raise ImportError(
"Could not import tensorflow_text python package. "
"Please install it with `pip install tensorflow_text``."
)
self.embed = tensorflow_hub.load(self.model_url)
model_config = ConfigDict(
extra="forbid",
)
def embed_documents(self, texts: List[str]) -> List[List[float]]:
"""Compute doc embeddings using a TensorflowHub embedding model.
Args:
texts: The list of texts to embed.
Returns:
List of embeddings, one for each text.
"""
texts = list(map(lambda x: x.replace("\n", " "), texts))
embeddings = self.embed(texts).numpy()
return embeddings.tolist()
def embed_query(self, text: str) -> List[float]:
"""Compute query embeddings using a TensorflowHub embedding model.
Args:
text: The text to embed.
Returns:
Embeddings for the text.
"""
text = text.replace("\n", " ")
embedding = self.embed([text]).numpy()[0]
return embedding.tolist()