Embeddings

ragbits.core.embeddings.Embeddings

Bases: ABC

Abstract client for communication with embedding models.

embed_text abstractmethod async

embed_text(data: list[str]) -> list[list[float]]

Creates embeddings for the given strings.

PARAMETER	DESCRIPTION
data	List of strings to get embeddings for. TYPE: list[str]

RETURNS	DESCRIPTION
list[list[float]]	List of embeddings for the given strings.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

@abstractmethod
async def embed_text(self, data: list[str]) -> list[list[float]]:
    """
    Creates embeddings for the given strings.

    Args:
        data: List of strings to get embeddings for.

    Returns:
        List of embeddings for the given strings.
    """

image_support

image_support() -> bool

Check if the model supports image embeddings.

RETURNS	DESCRIPTION
bool	True if the model supports image embeddings, False otherwise.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

def image_support(self) -> bool:  # noqa: PLR6301
    """
    Check if the model supports image embeddings.

    Returns:
        True if the model supports image embeddings, False otherwise.
    """
    return False

embed_image async

embed_image(images: list[bytes]) -> list[list[float]]

Creates embeddings for the given images.

PARAMETER	DESCRIPTION
images	List of images to get embeddings for. TYPE: list[bytes]

RETURNS	DESCRIPTION
list[list[float]]	List of embeddings for the given images.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

async def embed_image(self, images: list[bytes]) -> list[list[float]]:
    """
    Creates embeddings for the given images.

    Args:
        images: List of images to get embeddings for.

    Returns:
        List of embeddings for the given images.
    """
    raise NotImplementedError("Image embeddings are not supported by this model.")

ragbits.core.embeddings.local.LocalEmbeddings

LocalEmbeddings(model_name: str, api_key: str | None = None)

Bases: Embeddings

Class for interaction with any encoder available in HuggingFace.

Constructs a new local LLM instance.

PARAMETER	DESCRIPTION
model_name	Name of the model to use. TYPE: str
api_key	The API key for Hugging Face authentication. TYPE: str | None DEFAULT: None

RAISES	DESCRIPTION
ImportError	If the 'local' extra requirements are not installed.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/local.py

def __init__(
    self,
    model_name: str,
    api_key: str | None = None,
) -> None:
    """Constructs a new local LLM instance.

    Args:
        model_name: Name of the model to use.
        api_key: The API key for Hugging Face authentication.

    Raises:
        ImportError: If the 'local' extra requirements are not installed.
    """
    if not HAS_LOCAL_EMBEDDINGS:
        raise ImportError("You need to install the 'local' extra requirements to use local embeddings models")

    super().__init__()

    self.hf_api_key = api_key
    self.model_name = model_name

    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    self.model = AutoModel.from_pretrained(self.model_name, token=self.hf_api_key).to(self.device)
    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, token=self.hf_api_key)

hf_api_key instance-attribute

hf_api_key = api_key

model_name instance-attribute

model_name = model_name

device instance-attribute

device = device('cuda' if is_available() else 'cpu')

model instance-attribute

model = to(device)

tokenizer instance-attribute

tokenizer = from_pretrained(model_name, token=hf_api_key)

image_support

image_support() -> bool

Check if the model supports image embeddings.

RETURNS	DESCRIPTION
bool	True if the model supports image embeddings, False otherwise.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

def image_support(self) -> bool:  # noqa: PLR6301
    """
    Check if the model supports image embeddings.

    Returns:
        True if the model supports image embeddings, False otherwise.
    """
    return False

embed_image async

embed_image(images: list[bytes]) -> list[list[float]]

Creates embeddings for the given images.

PARAMETER	DESCRIPTION
images	List of images to get embeddings for. TYPE: list[bytes]

RETURNS	DESCRIPTION
list[list[float]]	List of embeddings for the given images.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

async def embed_image(self, images: list[bytes]) -> list[list[float]]:
    """
    Creates embeddings for the given images.

    Args:
        images: List of images to get embeddings for.

    Returns:
        List of embeddings for the given images.
    """
    raise NotImplementedError("Image embeddings are not supported by this model.")

embed_text async

embed_text(data: list[str], batch_size: int = 1) -> list[list[float]]

Calls the appropriate encoder endpoint with the given data and options.

PARAMETER	DESCRIPTION
data	List of strings to get embeddings for. TYPE: list[str]
batch_size	Batch size. TYPE: int DEFAULT: 1

RETURNS	DESCRIPTION
list[list[float]]	List of embeddings for the given strings.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/local.py

async def embed_text(self, data: list[str], batch_size: int = 1) -> list[list[float]]:
    """Calls the appropriate encoder endpoint with the given data and options.

    Args:
        data: List of strings to get embeddings for.
        batch_size: Batch size.

    Returns:
        List of embeddings for the given strings.
    """
    embeddings = []
    for batch in self._batch(data, batch_size):
        batch_dict = self.tokenizer(
            batch,
            max_length=self.tokenizer.model_max_length,
            padding=True,
            truncation=True,
            return_tensors="pt",
        ).to(self.device)
        with torch.no_grad():
            outputs = self.model(**batch_dict)
            batch_embeddings = self._average_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
            batch_embeddings = F.normalize(batch_embeddings, p=2, dim=1)
        embeddings.extend(batch_embeddings.to("cpu").tolist())

    torch.cuda.empty_cache()
    return embeddings

ragbits.core.embeddings.litellm.LiteLLMEmbeddings

LiteLLMEmbeddings(model: str = 'text-embedding-3-small', options: dict | None = None, api_base: str | None = None, api_key: str | None = None, api_version: str | None = None)

Bases: Embeddings

Client for creating text embeddings using LiteLLM API.

Constructs the LiteLLMEmbeddingClient.

PARAMETER	DESCRIPTION
model	Name of the LiteLLM supported model to be used. Default is "text-embedding-3-small". TYPE: str DEFAULT: 'text-embedding-3-small'
options	Additional options to pass to the LiteLLM API. TYPE: dict | None DEFAULT: None
api_base	The API endpoint you want to call the model with. TYPE: str | None DEFAULT: None
api_key	API key to be used. API key to be used. If not specified, an environment variable will be used, for more information, follow the instructions for your specific vendor in the LiteLLM documentation. TYPE: str | None DEFAULT: None
api_version	The API version for the call. TYPE: str | None DEFAULT: None

RAISES	DESCRIPTION
ImportError	If the 'litellm' extra requirements are not installed.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/litellm.py

def __init__(
    self,
    model: str = "text-embedding-3-small",
    options: dict | None = None,
    api_base: str | None = None,
    api_key: str | None = None,
    api_version: str | None = None,
) -> None:
    """
    Constructs the LiteLLMEmbeddingClient.

    Args:
        model: Name of the [LiteLLM supported model](https://docs.litellm.ai/docs/embedding/supported_embedding)\
            to be used. Default is "text-embedding-3-small".
        options: Additional options to pass to the LiteLLM API.
        api_base: The API endpoint you want to call the model with.
        api_key: API key to be used. API key to be used. If not specified, an environment variable will be used,
            for more information, follow the instructions for your specific vendor in the\
            [LiteLLM documentation](https://docs.litellm.ai/docs/embedding/supported_embedding).
        api_version: The API version for the call.

    Raises:
        ImportError: If the 'litellm' extra requirements are not installed.
    """
    if not HAS_LITELLM:
        raise ImportError("You need to install the 'litellm' extra requirements to use LiteLLM embeddings models")

    super().__init__()
    self.model = model
    self.options = options or {}
    self.api_base = api_base
    self.api_key = api_key
    self.api_version = api_version

model instance-attribute

model = model

options instance-attribute

options = options or {}

api_base instance-attribute

api_base = api_base

api_key instance-attribute

api_key = api_key

api_version instance-attribute

api_version = api_version

image_support

image_support() -> bool

Check if the model supports image embeddings.

RETURNS	DESCRIPTION
bool	True if the model supports image embeddings, False otherwise.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

def image_support(self) -> bool:  # noqa: PLR6301
    """
    Check if the model supports image embeddings.

    Returns:
        True if the model supports image embeddings, False otherwise.
    """
    return False

embed_image async

embed_image(images: list[bytes]) -> list[list[float]]

Creates embeddings for the given images.

PARAMETER	DESCRIPTION
images	List of images to get embeddings for. TYPE: list[bytes]

RETURNS	DESCRIPTION
list[list[float]]	List of embeddings for the given images.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/base.py

async def embed_image(self, images: list[bytes]) -> list[list[float]]:
    """
    Creates embeddings for the given images.

    Args:
        images: List of images to get embeddings for.

    Returns:
        List of embeddings for the given images.
    """
    raise NotImplementedError("Image embeddings are not supported by this model.")

embed_text async

embed_text(data: list[str]) -> list[list[float]]

Creates embeddings for the given strings.

PARAMETER	DESCRIPTION
data	List of strings to get embeddings for. TYPE: list[str]

RETURNS	DESCRIPTION
list[list[float]]	List of embeddings for the given strings.

RAISES	DESCRIPTION
EmbeddingConnectionError	If there is a connection error with the embedding API.
EmbeddingEmptyResponseError	If the embedding API returns an empty response.
EmbeddingStatusError	If the embedding API returns an error status code.
EmbeddingResponseError	If the embedding API response is invalid.

Source code in packages/ragbits-core/src/ragbits/core/embeddings/litellm.py

async def embed_text(self, data: list[str]) -> list[list[float]]:
    """
    Creates embeddings for the given strings.

    Args:
        data: List of strings to get embeddings for.

    Returns:
        List of embeddings for the given strings.

    Raises:
        EmbeddingConnectionError: If there is a connection error with the embedding API.
        EmbeddingEmptyResponseError: If the embedding API returns an empty response.
        EmbeddingStatusError: If the embedding API returns an error status code.
        EmbeddingResponseError: If the embedding API response is invalid.
    """
    with trace(
        data=data,
        model=self.model,
        api_base=self.api_base,
        api_version=self.api_version,
        options=self.options,
    ) as outputs:
        try:
            response = await litellm.aembedding(
                input=data,
                model=self.model,
                api_base=self.api_base,
                api_key=self.api_key,
                api_version=self.api_version,
                **self.options,
            )
        except litellm.openai.APIConnectionError as exc:
            raise EmbeddingConnectionError() from exc
        except litellm.openai.APIStatusError as exc:
            raise EmbeddingStatusError(exc.message, exc.status_code) from exc
        except litellm.openai.APIResponseValidationError as exc:
            raise EmbeddingResponseError() from exc

        if not response.data:
            raise EmbeddingEmptyResponseError()

        outputs.embeddings = [embedding["embedding"] for embedding in response.data]
        if response.usage:
            outputs.completion_tokens = response.usage.completion_tokens
            outputs.prompt_tokens = response.usage.prompt_tokens
            outputs.total_tokens = response.usage.total_tokens

    return outputs.embeddings