Creating your own embedding function

  1. from chromadb.api.types import (
  2.     Documents,
  3.     EmbeddingFunction,
  4.     Embeddings
  5. )
  8. class MyCustomEmbeddingFunction(EmbeddingFunction[Documents]):
  9.     def __init__(
  10.             self,
  11.             my_ef_param: str
  12.     ):
  13.         """Initialize the embedding function."""
  15.     def __call__(self, input: Documents) -> Embeddings:
  16.         """Embed the input documents."""
  17.         return self._my_ef(input)

Now let’s break the above down.

First you create a class that inherits from EmbeddingFunction[Documents]. The Documents type is a list of Document objects. Each Document object has a text attribute that contains the text of the document. Chroma also supports multi-modal

Example Implementation

Below is an implementation of an embedding function that works with transformers models.

Note

This example requires the transformers and torch python packages. You can install them with pip install transformers torch.

By default, all transformers models on HF are supported are also supported by the sentence-transformers package. For which Chroma provides out of the box support.

  1. import importlib
  2. from typing import Optional, cast
  4. import numpy as np
  5. import numpy.typing as npt
  6. from chromadb.api.types import EmbeddingFunction, Documents, Embeddings
  9. class TransformerEmbeddingFunction(EmbeddingFunction[Documents]):
  10.     def __init__(
  11.             self,
  12.             model_name: str = "dbmdz/bert-base-turkish-cased",
  13.             cache_dir: Optional[str] = None,
  14.     ):
  15.         try:
  16.             from transformers import AutoModel, AutoTokenizer
  18.             self._torch = importlib.import_module("torch")
  19.             self._tokenizer = AutoTokenizer.from_pretrained(model_name)
  20.             self._model = AutoModel.from_pretrained(model_name, cache_dir=cache_dir)
  21.         except ImportError:
  22.             raise ValueError(
  23.                 "The transformers and/or pytorch python package is not installed. Please install it with "
  24.                 "`pip install transformers` or `pip install torch`"
  25.             )
  27.     @staticmethod
  28.     def _normalize(vector: npt.NDArray) -> npt.NDArray:
  29.         """Normalizes a vector to unit length using L2 norm."""
  30.         norm = np.linalg.norm(vector)
  31.         if norm == 0:
  32.             return vector
  33.         return vector / norm
  35.     def __call__(self, input: Documents) -> Embeddings:
  36.         inputs = self._tokenizer(
  37.             input, padding=True, truncation=True, return_tensors="pt"
  38.         )
  39.         with self._torch.no_grad():
  40.             outputs = self._model(**inputs)
  41.         embeddings = outputs.last_hidden_state.mean(dim=1)  # mean pooling
  42.         return [e.tolist() for e in self._normalize(embeddings)]

May 18, 2024