code: import logging
  from typing import Any, List, Optional
  
  from llama_index.legacy.bridge.pydantic import Field, PrivateAttr
  from llama_index.legacy.core.embeddings.base import (
      DEFAULT_EMBED_BATCH_SIZE,
      BaseEmbedding,
      Embedding,
  )
  
  logger = logging.getLogger(__name__)
  
  
  # For bge models that Gradient AI provides, it is suggested to add the instruction for retrieval.
  # Reference: https://huggingface.co/BAAI/bge-large-en-v1.5#model-list
  QUERY_INSTRUCTION_FOR_RETRIEVAL = (
      "Represent this sentence for searching relevant passages:"
  )
  
  GRADIENT_EMBED_BATCH_SIZE: int = 32_768
  
  
  class GradientEmbedding(BaseEmbedding):
      """GradientAI embedding models.
  
      This class provides an interface to generate embeddings using a model
      deployed in Gradient AI. At the initialization it requires a model_id
      of the model deployed in the cluster.
  
      Note:
          Requires `gradientai` package to be available in the PYTHONPATH. It can be installed with
          `pip install gradientai`.
      """
  
      embed_batch_size: int = Field(default=GRADIENT_EMBED_BATCH_SIZE, gt=0)
  
      _gradient: Any = PrivateAttr()
      _model: Any = PrivateAttr()
  
      @classmethod
      def class_name(cls) -> str:
          return "GradientEmbedding"
  
      def __init__(
          self,
          *,
          embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
          gradient_model_slug: str,
          gradient_access_token: Optional[str] = None,
          gradient_workspace_id: Optional[str] = None,
          gradient_host: Optional[str] = None,
          **kwargs: Any,
      ):
          """Initializes the GradientEmbedding class.
  
          During the initialization the `gradientai` package is imported. Using the access token,
          workspace id and the slug of the model, the model is fetched from Gradient AI and prepared to use.
  
          Args:
              embed_batch_size (int, optional): The batch size for embedding generation. Defaults to 10,
                  must be > 0 and <= 100.
              gradient_model_slug (str): The model slug of the model in the Gradient AI account.
              gradient_access_token (str, optional): The access token of the Gradient AI account, if
                  `None` read from the environment variable `GRADIENT_ACCESS_TOKEN`.
              gradient_workspace_id (str, optional): The workspace ID of the Gradient AI account, if `None`
                  read from the environment variable `GRADIENT_WORKSPACE_ID`.
              gradient_host (str, optional): The host of the Gradient AI API. Defaults to None, which
                means the default host is used.
  
          Raises:
              ImportError: If the `gradientai` package is not available in the PYTHONPATH.
              ValueError: If the model cannot be fetched from Gradient AI.
          """
          if embed_batch_size <= 0:
              raise ValueError(f"Embed batch size {embed_batch_size}  must be > 0.")
  
          try:
              import gradientai
          except ImportError:
              raise ImportError("GradientEmbedding requires `pip install gradientai`.")
  
          self._gradient = gradientai.Gradient(
              access_token=gradient_access_token,
              workspace_id=gradient_workspace_id,
              host=gradient_host,
          )
  
          try:
              self._model = self._gradient.get_embeddings_model(slug=gradient_model_slug)
          except gradientai.openapi.client.exceptions.UnauthorizedException as e:
              logger.error(f"Error while loading model {gradient_model_slug}.")
              self._gradient.close()
              raise ValueError("Unable to fetch the requested embeddings model") from e
  
          super().__init__(
              embed_batch_size=embed_batch_size, model_name=gradient_model_slug, **kwargs
          )
  
      async def _aget_text_embeddings(self, texts: List[str]) -> List[Embedding]:
          """
          Embed the input sequence of text asynchronously.
          """
          inputs = [{"input": text} for text in texts]
  
          result = await self._model.aembed(inputs=inputs).embeddings
  
          return [e.embedding for e in result]
  
      def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]:
          """
          Embed the input sequence of text.
          """
          inputs = [{"input": text} for text in texts]
  
          result = self._model.embed(inputs=inputs).embeddings
  
          return [e.embedding for e in result]
  
      def _get_text_embedding(self, text: str) -> Embedding:
          """Alias for _get_text_embeddings() with single text input."""
          return self._get_text_embeddings([text])[0]
  
      async def _aget_text_embedding(self, text: str) -> Embedding:
          """Alias for _aget_text_embeddings() with single text input."""
          embedding = await self._aget_text_embeddings([text])
          return embedding[0]
  
      async def _aget_query_embedding(self, query: str) -> Embedding:
          embedding = await self._aget_text_embeddings(
              [f"{QUERY_INSTRUCTION_FOR_RETRIEVAL} {query}"]
          )
          return embedding[0]
  
      def _get_query_embedding(self, query: str) -> Embedding:
          return self._get_text_embeddings(
              [f"{QUERY_INSTRUCTION_FOR_RETRIEVAL} {query}"]
          )[0]
  
apis: llama_index.legacy.bridge.pydantic.PrivateAttr llama_index.legacy.bridge.pydantic.Field
extract_api: [((251, 278), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (268, 278), False, 'import logging\n'), ((1040, 1086), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'GRADIENT_EMBED_BATCH_SIZE', 'gt': '(0)'}), '(default=GRADIENT_EMBED_BATCH_SIZE, gt=0)\n', (1045, 1086), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1109, 1122), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (1120, 1122), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1141, 1154), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (1152, 1154), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((3068, 3184), 'gradientai.Gradient', 'gradientai.Gradient', ([], {'access_token': 'gradient_access_token', 'workspace_id': 'gradient_workspace_id', 'host': 'gradient_host'}), '(access_token=gradient_access_token, workspace_id=\n    gradient_workspace_id, host=gradient_host)\n', (3087, 3184), False, 'import gradientai\n')]