Source code for langchain_community.embeddings.vertexai
import logging
import re
import string
import threading
from concurrent.futures import ThreadPoolExecutor, wait
from typing import Any, Dict, List, Literal, Optional, Tuple
from langchain_core.embeddings import Embeddings
from langchain_core.language_models.llms import create_base_retry_decorator
from langchain_core.pydantic_v1 import root_validator
from langchain_community.llms.vertexai import _VertexAICommon
from langchain_community.utilities.vertexai import raise_vertex_import_error
logger = logging.getLogger(__name__)
_MAX_TOKENS_PER_BATCH = 20000
_MAX_BATCH_SIZE = 250
_MIN_BATCH_SIZE = 5
[docs]class VertexAIEmbeddings(_VertexAICommon, Embeddings):
    """Google Cloud VertexAI embedding models."""
    # Instance context
    instance: Dict[str, Any] = {}  #: :meta private:
    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validates that the python package exists in environment."""
        cls._try_init_vertexai(values)
        try:
            from vertexai.language_models import TextEmbeddingModel
            values["client"] = TextEmbeddingModel.from_pretrained(values["model_name"])
        except ImportError:
            raise_vertex_import_error()
        return values
    def __init__(
        self,
        project: Optional[str] = None,
        location: str = "us-central1",
        request_parallelism: int = 5,
        max_retries: int = 6,
        model_name: str = "textembedding-gecko",
        credentials: Optional[Any] = None,
        **kwargs: Any,
    ):
        """Initialize the sentence_transformer."""
        super().__init__(
            project=project,
            location=location,