chore: import upstream snapshot with attribution

unstructured/embed/README.md

@@ -0,0 +1,6 @@
+# Embed
+![Project unmaintained](https://img.shields.io/badge/project-unmaintained-red.svg)
+
+Project has been moved to: [Unstructured Ingest](https://github.com/Unstructured-IO/unstructured-ingest)
+
+This python module will be removed from this repo in the near future.

unstructured/embed/__init__.py

@@ -0,0 +1,27 @@
+import warnings
+
+from unstructured.embed.bedrock import BedrockEmbeddingEncoder
+from unstructured.embed.huggingface import HuggingFaceEmbeddingEncoder
+from unstructured.embed.mixedbreadai import MixedbreadAIEmbeddingEncoder
+from unstructured.embed.octoai import OctoAIEmbeddingEncoder
+from unstructured.embed.openai import OpenAIEmbeddingEncoder
+from unstructured.embed.vertexai import VertexAIEmbeddingEncoder
+from unstructured.embed.voyageai import VoyageAIEmbeddingEncoder
+
+EMBEDDING_PROVIDER_TO_CLASS_MAP = {
+    "langchain-openai": OpenAIEmbeddingEncoder,
+    "langchain-huggingface": HuggingFaceEmbeddingEncoder,
+    "langchain-aws-bedrock": BedrockEmbeddingEncoder,
+    "langchain-vertexai": VertexAIEmbeddingEncoder,
+    "voyageai": VoyageAIEmbeddingEncoder,
+    "mixedbread-ai": MixedbreadAIEmbeddingEncoder,
+    "octoai": OctoAIEmbeddingEncoder,
+}
+
+
+warnings.warn(
+    "unstructured.ingest will be removed in a future version. "
+    "Functionality moved to the unstructured-ingest project.",
+    DeprecationWarning,
+    stacklevel=2,
+)

unstructured/embed/bedrock.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List
+
+import numpy as np
+from pydantic import SecretStr
+
+from unstructured.documents.elements import (
+    Element,
+)
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import requires_dependencies
+
+if TYPE_CHECKING:
+    from langchain_community.embeddings import BedrockEmbeddings
+
+
+class BedrockEmbeddingConfig(EmbeddingConfig):
+    aws_access_key_id: SecretStr
+    aws_secret_access_key: SecretStr
+    region_name: str = "us-west-2"
+
+    @requires_dependencies(
+        ["boto3", "numpy", "langchain_community"],
+        extras="bedrock",
+    )
+    def get_client(self) -> "BedrockEmbeddings":
+        # delay import only when needed
+        import boto3
+        from langchain_community.embeddings import BedrockEmbeddings
+
+        bedrock_runtime = boto3.client(
+            service_name="bedrock-runtime",
+            aws_access_key_id=self.aws_access_key_id.get_secret_value(),
+            aws_secret_access_key=self.aws_secret_access_key.get_secret_value(),
+            region_name=self.region_name,
+        )
+
+        bedrock_client = BedrockEmbeddings(client=bedrock_runtime)
+        return bedrock_client
+
+
+@dataclass
+class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: BedrockEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> List[float]:
+        return self.embed_query(query="Q")
+
+    def __post_init__(self):
+        self.initialize()
+
+    def num_of_dimensions(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def is_unit_vector(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def embed_query(self, query):
+        bedrock_client = self.config.get_client()
+        return np.array(bedrock_client.embed_query(query))
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        bedrock_client = self.config.get_client()
+        embeddings = bedrock_client.embed_documents([str(e) for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+        for i, element in enumerate(elements):
+            element.embeddings = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements

unstructured/embed/huggingface.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+from pydantic import Field
+
+from unstructured.documents.elements import (
+    Element,
+)
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import requires_dependencies
+
+if TYPE_CHECKING:
+    from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+
+
+class HuggingFaceEmbeddingConfig(EmbeddingConfig):
+    model_name: Optional[str] = Field(default="sentence-transformers/all-MiniLM-L6-v2")
+    model_kwargs: Optional[dict] = Field(default_factory=lambda: {"device": "cpu"})
+    encode_kwargs: Optional[dict] = Field(default_factory=lambda: {"normalize_embeddings": False})
+    cache_folder: Optional[dict] = Field(default=None)
+
+    @requires_dependencies(
+        ["langchain_huggingface"],
+        extras="embed-huggingface",
+    )
+    def get_client(self) -> "HuggingFaceEmbeddings":
+        """Creates a langchain Huggingface python client to embed elements."""
+        from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+
+        client = HuggingFaceEmbeddings(**self.dict())
+        return client
+
+
+@dataclass
+class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: HuggingFaceEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> List[float]:
+        return self.embed_query(query="Q")
+
+    def num_of_dimensions(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def is_unit_vector(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def embed_query(self, query):
+        client = self.config.get_client()
+        return client.embed_query(str(query))
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        client = self.config.get_client()
+        embeddings = client.embed_documents([str(e) for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+
+        for i, element in enumerate(elements):
+            element.embeddings = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements

unstructured/embed/interfaces.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Tuple
+
+from pydantic import BaseModel
+
+from unstructured.documents.elements import Element
+
+
+class EmbeddingConfig(BaseModel):
+    pass
+
+
+@dataclass
+class BaseEmbeddingEncoder(ABC):
+    config: EmbeddingConfig
+
+    @abstractmethod
+    def initialize(self):
+        """Initializes the embedding encoder class. Should also validate the instance
+        is properly configured: e.g., embed a single a element"""
+
+    @property
+    @abstractmethod
+    def num_of_dimensions(self) -> Tuple[int]:
+        """Number of dimensions for the embedding vector."""
+
+    @property
+    @abstractmethod
+    def is_unit_vector(self) -> bool:
+        """Denotes if the embedding vector is a unit vector."""
+
+    @abstractmethod
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        pass
+
+    @abstractmethod
+    def embed_query(self, query: str) -> List[float]:
+        pass

unstructured/embed/mixedbreadai.py

@@ -0,0 +1,178 @@
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+from pydantic import Field, SecretStr
+
+from unstructured.documents.elements import Element
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import requires_dependencies
+
+USER_AGENT = "@mixedbread-ai/unstructured"
+BATCH_SIZE = 128
+TIMEOUT = 60
+MAX_RETRIES = 3
+ENCODING_FORMAT = "float"
+TRUNCATION_STRATEGY = "end"
+
+
+if TYPE_CHECKING:
+    from mixedbread_ai.client import MixedbreadAI
+    from mixedbread_ai.core import RequestOptions
+
+
+class MixedbreadAIEmbeddingConfig(EmbeddingConfig):
+    """
+    Configuration class for Mixedbread AI Embedding Encoder.
+
+    Attributes:
+        api_key (str): API key for accessing Mixedbread AI..
+        model_name (str): Name of the model to use for embeddings.
+    """
+
+    api_key: SecretStr = Field(
+        default_factory=lambda: SecretStr(os.environ.get("MXBAI_API_KEY")),
+    )
+
+    model_name: str = Field(
+        default="mixedbread-ai/mxbai-embed-large-v1",
+    )
+
+    @requires_dependencies(
+        ["mixedbread_ai"],
+        extras="embed-mixedbreadai",
+    )
+    def get_client(self) -> "MixedbreadAI":
+        """
+        Create the Mixedbread AI client.
+
+        Returns:
+            MixedbreadAI: Initialized client.
+        """
+        from mixedbread_ai.client import MixedbreadAI
+
+        return MixedbreadAI(
+            api_key=self.api_key.get_secret_value(),
+        )
+
+
+@dataclass
+class MixedbreadAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    """
+    Embedding encoder for Mixedbread AI.
+
+    Attributes:
+        config (MixedbreadAIEmbeddingConfig): Configuration for the embedding encoder.
+    """
+
+    config: MixedbreadAIEmbeddingConfig
+
+    _exemplary_embedding: Optional[List[float]] = field(init=False, default=None)
+    _request_options: Optional["RequestOptions"] = field(init=False, default=None)
+
+    def get_exemplary_embedding(self) -> List[float]:
+        """Get an exemplary embedding to determine dimensions and unit vector status."""
+        return self._embed(["Q"])[0]
+
+    def initialize(self):
+        if self.config.api_key is None:
+            raise ValueError(
+                "The Mixedbread AI API key must be specified."
+                + "You either pass it in the constructor using 'api_key'"
+                + "or via the 'MXBAI_API_KEY' environment variable."
+            )
+
+        from mixedbread_ai.core import RequestOptions
+
+        self._request_options = RequestOptions(
+            max_retries=MAX_RETRIES,
+            timeout_in_seconds=TIMEOUT,
+            additional_headers={"User-Agent": USER_AGENT},
+        )
+
+    @property
+    def num_of_dimensions(self):
+        """Get the number of dimensions for the embeddings."""
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    @property
+    def is_unit_vector(self) -> bool:
+        """Check if the embedding is a unit vector."""
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def _embed(self, texts: List[str]) -> List[List[float]]:
+        """
+        Embed a list of texts using the Mixedbread AI API.
+
+        Args:
+            texts (List[str]): List of texts to embed.
+
+        Returns:
+            List[List[float]]: List of embeddings.
+        """
+        batch_size = BATCH_SIZE
+        batch_itr = range(0, len(texts), batch_size)
+
+        responses = []
+        client = self.config.get_client()
+        for i in batch_itr:
+            batch = texts[i : i + batch_size]
+            response = client.embeddings(
+                model=self.config.model_name,
+                normalized=True,
+                encoding_format=ENCODING_FORMAT,
+                truncation_strategy=TRUNCATION_STRATEGY,
+                request_options=self._request_options,
+                input=batch,
+            )
+            responses.append(response)
+        return [item.embedding for response in responses for item in response.data]
+
+    @staticmethod
+    def _add_embeddings_to_elements(
+        elements: List[Element], embeddings: List[List[float]]
+    ) -> List[Element]:
+        """
+        Add embeddings to elements.
+
+        Args:
+            elements (List[Element]): List of elements.
+            embeddings (List[List[float]]): List of embeddings.
+
+        Returns:
+            List[Element]: Elements with embeddings added.
+        """
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+        for i, element in enumerate(elements):
+            element.embeddings = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        """
+        Embed a list of document elements.
+
+        Args:
+            elements (List[Element]): List of document elements.
+
+        Returns:
+            List[Element]: Elements with embeddings.
+        """
+        embeddings = self._embed([str(e) for e in elements])
+        return self._add_embeddings_to_elements(elements, embeddings)
+
+    def embed_query(self, query: str) -> List[float]:
+        """
+        Embed a query string.
+
+        Args:
+            query (str): Query string to embed.
+
+        Returns:
+            List[float]: Embedding of the query.
+        """
+        return self._embed([query])[0]

unstructured/embed/octoai.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+from pydantic import Field, SecretStr
+
+from unstructured.documents.elements import (
+    Element,
+)
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import requires_dependencies
+
+if TYPE_CHECKING:
+    from openai import OpenAI
+
+
+class OctoAiEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    model_name: str = Field(default="thenlper/gte-large")
+    base_url: str = Field(default="https://text.octoai.run/v1")
+
+    @requires_dependencies(
+        ["openai", "tiktoken"],
+        extras="embed-octoai",
+    )
+    def get_client(self) -> "OpenAI":
+        """Creates an OpenAI python client to embed elements. Uses the OpenAI SDK."""
+        from openai import OpenAI
+
+        return OpenAI(api_key=self.api_key.get_secret_value(), base_url=self.base_url)
+
+
+@dataclass
+class OctoAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: OctoAiEmbeddingConfig
+    # Uses the OpenAI SDK
+    _exemplary_embedding: Optional[List[float]] = field(init=False, default=None)
+
+    def get_exemplary_embedding(self) -> List[float]:
+        return self.embed_query("Q")
+
+    def initialize(self):
+        pass
+
+    def num_of_dimensions(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def is_unit_vector(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def embed_query(self, query):
+        client = self.config.get_client()
+        response = client.embeddings.create(input=str(query), model=self.config.model_name)
+        return response.data[0].embedding
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        embeddings = [self.embed_query(e) for e in elements]
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+        for i, element in enumerate(elements):
+            element.embeddings = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements

unstructured/embed/openai.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List
+
+import numpy as np
+from pydantic import Field, SecretStr
+
+from unstructured.documents.elements import (
+    Element,
+)
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import requires_dependencies
+
+if TYPE_CHECKING:
+    from langchain_openai.embeddings import OpenAIEmbeddings
+
+
+class OpenAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    model_name: str = Field(default="text-embedding-ada-002")
+
+    @requires_dependencies(["langchain_openai"], extras="openai")
+    def get_client(self) -> "OpenAIEmbeddings":
+        """Creates a langchain OpenAI python client to embed elements."""
+        from langchain_openai import OpenAIEmbeddings
+
+        openai_client = OpenAIEmbeddings(
+            openai_api_key=self.api_key.get_secret_value(),
+            model=self.model_name,  # type: ignore
+        )
+        return openai_client
+
+
+@dataclass
+class OpenAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: OpenAIEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> List[float]:
+        return self.embed_query(query="Q")
+
+    def initialize(self):
+        pass
+
+    def num_of_dimensions(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def is_unit_vector(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def embed_query(self, query):
+        client = self.config.get_client()
+        return client.embed_query(str(query))
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        client = self.config.get_client()
+        embeddings = client.embed_documents([str(e) for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+        for i, element in enumerate(elements):
+            element.embeddings = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements

unstructured/embed/vertexai.py

@@ -0,0 +1,76 @@
+# type: ignore
+import json
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+from pydantic import Field, SecretStr
+
+from unstructured.documents.elements import (
+    Element,
+)
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import FileHandler, requires_dependencies
+
+if TYPE_CHECKING:
+    from langchain_google_vertexai import VertexAIEmbeddings
+
+
+class VertexAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    model_name: Optional[str] = Field(default="textembedding-gecko@001")
+
+    def register_application_credentials(self):
+        application_credentials_path = os.path.join("/tmp", "google-vertex-app-credentials.json")
+        credentials_file = FileHandler(application_credentials_path)
+        credentials_file.write_file(json.dumps(json.loads(self.api_key.get_secret_value())))
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = application_credentials_path
+
+    @requires_dependencies(
+        ["langchain", "langchain_google_vertexai"],
+        extras="embed-vertexai",
+    )
+    def get_client(self) -> "VertexAIEmbeddings":
+        """Creates a Langchain VertexAI python client to embed elements."""
+        from langchain_google_vertexai import VertexAIEmbeddings
+
+        self.register_application_credentials()
+        vertexai_client = VertexAIEmbeddings(model_name=self.model_name)
+        return vertexai_client
+
+
+@dataclass
+class VertexAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: VertexAIEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> List[float]:
+        return self.embed_query(query="A sample query.")
+
+    def initialize(self):
+        pass
+
+    def num_of_dimensions(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    def is_unit_vector(self):
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def embed_query(self, query):
+        client = self.config.get_client()
+        result = client.embed_query(str(query))
+        return result
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        client = self.config.get_client()
+        embeddings = client.embed_documents([str(e) for e in elements])
+        elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
+        return elements_with_embeddings
+
+    def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
+        assert len(elements) == len(embeddings)
+        for element, embedding in zip(elements, embeddings):
+            element.embeddings = embedding
+        return elements

unstructured/embed/voyageai.py

@@ -0,0 +1,237 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Iterable, List, Optional
+
+import numpy as np
+from pydantic import Field, SecretStr
+
+from unstructured.documents.elements import Element
+from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
+from unstructured.utils import requires_dependencies
+
+if TYPE_CHECKING:
+    from voyageai import Client
+
+# Token limits for different VoyageAI models
+VOYAGE_TOTAL_TOKEN_LIMITS = {
+    "voyage-context-3": 32_000,
+    "voyage-3.5-lite": 1_000_000,
+    "voyage-3.5": 320_000,
+    "voyage-2": 320_000,
+    "voyage-02": 320_000,
+    "voyage-3-large": 120_000,
+    "voyage-code-3": 120_000,
+    "voyage-large-2-instruct": 120_000,
+    "voyage-finance-2": 120_000,
+    "voyage-multilingual-2": 120_000,
+    "voyage-law-2": 120_000,
+    "voyage-large-2": 120_000,
+    "voyage-3": 120_000,
+    "voyage-3-lite": 120_000,
+    "voyage-code-2": 120_000,
+    "voyage-3-m-exp": 120_000,
+    "voyage-multimodal-3": 120_000,
+}
+
+# Batch size for embedding requests (max documents per batch)
+MAX_BATCH_SIZE = 1000
+
+
+class VoyageAIEmbeddingConfig(EmbeddingConfig):
+    api_key: SecretStr
+    model_name: str
+    show_progress_bar: bool = False
+    batch_size: Optional[int] = Field(default=None)
+    truncation: Optional[bool] = Field(default=None)
+    output_dimension: Optional[int] = Field(default=None)
+
+    @requires_dependencies(
+        ["voyageai"],
+        extras="embed-voyageai",
+    )
+    def get_client(self) -> "Client":
+        """Creates a VoyageAI python client to embed elements."""
+        from voyageai import Client
+
+        return Client(
+            api_key=self.api_key.get_secret_value(),
+        )
+
+    def get_token_limit(self) -> int:
+        """Get the token limit for the current model."""
+        return VOYAGE_TOTAL_TOKEN_LIMITS.get(self.model_name, 120_000)
+
+
+@dataclass
+class VoyageAIEmbeddingEncoder(BaseEmbeddingEncoder):
+    config: VoyageAIEmbeddingConfig
+
+    def get_exemplary_embedding(self) -> List[float]:
+        return self.embed_query(query="A sample query.")
+
+    def initialize(self):
+        pass
+
+    @property
+    def num_of_dimensions(self) -> tuple[int, ...]:
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.shape(exemplary_embedding)
+
+    @property
+    def is_unit_vector(self) -> bool:
+        exemplary_embedding = self.get_exemplary_embedding()
+        return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
+
+    def _is_context_model(self) -> bool:
+        """Check if the model is a contextualized embedding model."""
+        return "context" in self.config.model_name
+
+    def _build_batches(self, texts: List[str], client: "Client") -> Iterable[List[str]]:
+        """
+        Generate batches of texts based on token limits.
+
+        Args:
+            texts: List of texts to batch.
+            client: VoyageAI client instance to use for tokenization.
+
+        Yields:
+            Batches of texts as lists.
+        """
+        if not texts:
+            return
+
+        max_tokens_per_batch = self.config.get_token_limit()
+        current_batch: List[str] = []
+        current_batch_tokens = 0
+
+        # Tokenize all texts in one API call
+        all_token_lists = client.tokenize(texts, model=self.config.model_name)
+        token_counts = [len(tokens) for tokens in all_token_lists]
+
+        for i, text in enumerate(texts):
+            n_tokens = token_counts[i]
+
+            # Check if adding this text would exceed limits
+            if current_batch and (
+                len(current_batch) >= MAX_BATCH_SIZE
+                or (current_batch_tokens + n_tokens > max_tokens_per_batch)
+            ):
+                # Yield the current batch and start a new one
+                yield current_batch
+                current_batch = []
+                current_batch_tokens = 0
+
+            current_batch.append(text)
+            current_batch_tokens += n_tokens
+
+        # Yield the last batch (always has at least one text)
+        if current_batch:
+            yield current_batch
+
+    def _embed_batch(
+        self, batch: List[str], client: "Client", input_type: str = "document"
+    ) -> List[List[float]]:
+        """
+        Embed a batch of texts using the appropriate method for the model.
+
+        Args:
+            batch: List of texts to embed.
+            client: VoyageAI client instance to use for embedding.
+            input_type: Type of input ("document" or "query").
+
+        Returns:
+            List of embedding vectors.
+        """
+        if self._is_context_model():
+            result = client.contextualized_embed(
+                inputs=[batch],
+                model=self.config.model_name,
+                input_type=input_type,
+                output_dimension=self.config.output_dimension,
+            )
+            return [list(emb) for emb in result.results[0].embeddings]
+        else:
+            result = client.embed(
+                texts=batch,
+                model=self.config.model_name,
+                input_type=input_type,
+                truncation=self.config.truncation,
+                output_dimension=self.config.output_dimension,
+            )
+            return [list(emb) for emb in result.embeddings]
+
+    def embed_documents(self, elements: List[Element]) -> List[Element]:
+        """
+        Embed documents with automatic batching based on token limits.
+
+        Args:
+            elements: List of elements to embed.
+
+        Returns:
+            List of elements with embeddings added.
+        """
+        if not elements:
+            return []
+
+        client = self.config.get_client()
+        texts = [str(e) for e in elements]
+        all_embeddings: List[List[float]] = []
+
+        # Process each batch
+        batches = list(self._build_batches(texts, client))
+
+        if self.config.show_progress_bar:
+            try:
+                from tqdm.auto import tqdm  # type: ignore
+
+                batches = tqdm(batches, desc="Embedding batches")
+            except ImportError as e:
+                raise ImportError(
+                    "Must have tqdm installed if `show_progress_bar` is set to True. "
+                    "Please install with `pip install tqdm`."
+                ) from e
+
+        for batch in batches:
+            batch_embeddings = self._embed_batch(batch, client, input_type="document")
+            all_embeddings.extend(batch_embeddings)
+
+        return self._add_embeddings_to_elements(elements, all_embeddings)
+
+    def embed_query(self, query: str) -> List[float]:
+        """
+        Embed a single query string.
+
+        Args:
+            query: Query string to embed.
+
+        Returns:
+            Embedding vector.
+        """
+        client = self.config.get_client()
+        batch_embeddings = self._embed_batch([query], client, input_type="query")
+        return batch_embeddings[0]
+
+    def count_tokens(self, texts: List[str]) -> List[int]:
+        """
+        Count tokens for the given texts.
+
+        Args:
+            texts: List of texts to count tokens for.
+
+        Returns:
+            List of token counts for each text.
+        """
+        if not texts:
+            return []
+
+        client = self.config.get_client()
+        token_lists = client.tokenize(texts, model=self.config.model_name)
+        return [len(token_list) for token_list in token_lists]
+
+    @staticmethod
+    def _add_embeddings_to_elements(elements, embeddings) -> List[Element]:
+        assert len(elements) == len(embeddings)
+        elements_w_embedding = []
+        for i, element in enumerate(elements):
+            element.embeddings = embeddings[i]
+            elements_w_embedding.append(element)
+        return elements