Add Hugging Face Embedder Support

[evanmschultz]

Overview

The project currently supports OpenAI, Gemini, and Voyage embedders, all of which require API keys and incur usage costs. This issue proposes adding support for local Hugging Face embedding models, which are open source and can run locally without API costs.

Requirements

• Add support for local Hugging Face embedding models
• Make dependencies optional to avoid bloating the main package
• Maintain compatibility with the existing EmbedderClient interface
• Support async operations matching the current client pattern

Implementation Plan

1. Add Optional Dependencies

Update pyproject.toml to include Hugging Face dependencies as an optional requirement.

2. Create HuggingFaceEmbedder Class

Create a new file graphiti_core/embedder/huggingface.py implementing:

• HuggingFaceEmbedderConfig - Configuration class for model selection and parameters
• HuggingFaceEmbedder - Implementation of EmbedderClient for Hugging Face models

3. Update Embedder Module Exports

Update graphiti_core/embedder/__init__.py to export the new classes.

4. Implementation Details

HuggingFaceEmbedderConfig

• model_name: String for model selection (default to a small, general-purpose model)
• embedding_dim: Match the required dimension from the base config
• device: Option to specify CPU or GPU (with fallback to CPU)
• model_kwargs: Additional parameters to pass to the model

HuggingFaceEmbedder

• Implement the required create method for generating embeddings
• Add skeleton for create_batch with NotImplementedError for now
• Handle both synchronous and asynchronous operations
• Add lazy loading of models to avoid unnecessary imports when not used

5. Testing

• Add unit tests for the new embedder in the test suite
• Test with various model sizes and on different input types
• Verify compatibility with the rest of the Graphiti system
• Ensure graceful fallback when GPU is not available

6. Documentation

• Update documentation to include information about the new embedder

Technical Considerations

• The implementation will use the sentence-transformers package for model loading and embedding generation
• Models will be loaded on demand and cached for reuse
• May require thread/process management for async operations with CPU-only setup
• Default to small models to minimize memory footprint

[prasmussen15]

Yes, we definitely need to add this to graphiti. Funnily enough we use an HF embedder in our deployment of graphiti so we can just port that over

[evanmschultz]

Considering you have it in deployment, I imagine you would not want me to do anything to add it. Let me know if I can do anything to help with this!

[look4pritam]

HuggingFace embedding is essential for local deployment. Consider adding this feature.

[jfjeschke]

Here's the code:

import logging
from collections.abc import Iterable
from typing import Any

try:
    import sentence_transformers
except ImportError:
    raise ImportError(
        'sentence-transformers is required for HuggingFaceEmbedder. '
        'Install it with: pip install sentence-transformers'
    ) from None

from pydantic import Field

from .client import EmbedderClient, EmbedderConfig

logger = logging.getLogger(__name__)

DEFAULT_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"


class HuggingFaceEmbedderConfig(EmbedderConfig):
    model_name: str = Field(default=DEFAULT_MODEL_NAME)
    cache_folder: str | None = None
    model_kwargs: dict[str, Any] = Field(default_factory=dict)
    encode_kwargs: dict[str, Any] = Field(default_factory=dict)
    multi_process: bool = False
    show_progress: bool = False


class HuggingFaceEmbedder(EmbedderClient):
    """
    Hugging Face SentenceTransformers Embedder Client
    """

    def __init__(
        self,
        config: HuggingFaceEmbedderConfig | None = None,
        client: sentence_transformers.SentenceTransformer | None = None,
    ):
        """
        Initialize the HuggingFaceEmbedder with the provided configuration.

        Args:
            config (HuggingFaceEmbedderConfig | None): The configuration for the HuggingFaceEmbedder.
            client (sentence_transformers.SentenceTransformer | None): An existing SentenceTransformer instance.
                If provided, this will be used instead of creating a new one from the config.
        """
        if config is None:
            config = HuggingFaceEmbedderConfig()

        self.config = config

        # Use provided client or initialize a new sentence transformer model
        if client is not None:
            self._client = client
        else:
            self._client = sentence_transformers.SentenceTransformer(
                self.config.model_name,
                cache_folder=self.config.cache_folder,
                **self.config.model_kwargs
            )

    async def create(
        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
    ) -> list[float]:
        """
        Create embeddings for the given input data using Hugging Face SentenceTransformers.

        Args:
            input_data: The input data to create embeddings for. Can be a string, list of strings,
                       or an iterable of integers or iterables of integers.

        Returns:
            A list of floats representing the embedding vector.
        """
        # Convert input to text if needed
        if isinstance(input_data, str):
            text = input_data.replace("\n", " ")
        elif isinstance(input_data, list) and all(isinstance(item, str) for item in input_data):
            # If it's a list of strings, take the first one for single embedding
            text = input_data[0].replace("\n", " ") if input_data else ""
        else:
            # Convert other types to string
            text = str(input_data).replace("\n", " ")

        # Generate embedding
        embedding = self._client.encode(
            text,
            show_progress_bar=self.config.show_progress,
            **self.config.encode_kwargs
        )

        if hasattr(embedding, 'tolist'):
            return embedding.tolist()
        elif isinstance(embedding, list):
            return embedding
        else:
            raise TypeError(f"Unexpected embedding type: {type(embedding)}")

    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
        """
        Create embeddings for a batch of input data using Hugging Face SentenceTransformers.

        Args:
            input_data_list: A list of strings to create embeddings for.

        Returns:
            A list of embedding vectors (each vector is a list of floats).
        """
        if not input_data_list:
            return []

        try:
            # Clean the texts
            texts = [text.replace("\n", " ") for text in input_data_list]

            # Generate embeddings
            if self.config.multi_process:
                pool = self._client.start_multi_process_pool()
                embeddings = self._client.encode_multi_process(texts, pool)
                sentence_transformers.SentenceTransformer.stop_multi_process_pool(pool)
            else:
                embeddings = self._client.encode(
                    texts,
                    show_progress_bar=self.config.show_progress,
                    **self.config.encode_kwargs
                )

            if isinstance(embeddings, list):
                # If it's already a list, check if it contains the right type
                if embeddings and hasattr(embeddings[0], 'tolist'):
                    return [emb.tolist() for emb in embeddings]
                return embeddings
            else:
                # Convert tensor/array to list
                return embeddings.tolist()

        except Exception as e:
            logger.error(f"Error in Hugging Face batch embedding: {e}")
            raise

You can either pass in a config or an existing SentenceTransformer instance

[jiwon-hae]

Hi all, Is there currently someone actively working on integrating Hugging Face embedder support into Graphiti? If the implementation is still open or you need additional support, I’d be interested in contributing to this feature. Please let me know how I can help!

[claude]

@evanmschultz Is this still an issue? Please confirm within 14 days or this issue will be closed.

[claude]

@evanmschultz Is this still relevant? Please confirm within 14 days or this issue will be closed.

[anentropic]

Ollama is supported so I guess there is already a round-about route to running local embedding model (It'd be nice not to have to run a separate server though)

Also I'd love to know which is the preferred embedding model to use?