Rename embedchain to mem0 and open sourcing code for long term memory (#1474)
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
Taranjeet Singh
co-authored by
Deshraj Yadav
Deshraj Yadav
parent
83e8c97295
commit
f842a92e25
@@ -0,0 +1,22 @@
|
||||
from typing import Optional
|
||||
|
||||
from langchain_community.embeddings import AzureOpenAIEmbeddings
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class AzureOpenAIEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
if self.config.model is None:
|
||||
self.config.model = "text-embedding-ada-002"
|
||||
|
||||
embeddings = AzureOpenAIEmbeddings(deployment=self.config.deployment_name)
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
|
||||
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.OPENAI.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,90 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Optional
|
||||
|
||||
from embedchain.config.embedder.base import BaseEmbedderConfig
|
||||
|
||||
try:
|
||||
from chromadb.api.types import Embeddable, EmbeddingFunction, Embeddings
|
||||
except RuntimeError:
|
||||
from embedchain.utils.misc import use_pysqlite3
|
||||
|
||||
use_pysqlite3()
|
||||
from chromadb.api.types import Embeddable, EmbeddingFunction, Embeddings
|
||||
|
||||
|
||||
class EmbeddingFunc(EmbeddingFunction):
|
||||
def __init__(self, embedding_fn: Callable[[list[str]], list[str]]):
|
||||
self.embedding_fn = embedding_fn
|
||||
|
||||
def __call__(self, input: Embeddable) -> Embeddings:
|
||||
return self.embedding_fn(input)
|
||||
|
||||
|
||||
class BaseEmbedder:
|
||||
"""
|
||||
Class that manages everything regarding embeddings. Including embedding function, loaders and chunkers.
|
||||
|
||||
Embedding functions and vector dimensions are set based on the child class you choose.
|
||||
To manually overwrite you can use this classes `set_...` methods.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
"""
|
||||
Initialize the embedder class.
|
||||
|
||||
:param config: embedder configuration option class, defaults to None
|
||||
:type config: Optional[BaseEmbedderConfig], optional
|
||||
"""
|
||||
if config is None:
|
||||
self.config = BaseEmbedderConfig()
|
||||
else:
|
||||
self.config = config
|
||||
self.vector_dimension: int
|
||||
|
||||
def set_embedding_fn(self, embedding_fn: Callable[[list[str]], list[str]]):
|
||||
"""
|
||||
Set or overwrite the embedding function to be used by the database to store and retrieve documents.
|
||||
|
||||
:param embedding_fn: Function to be used to generate embeddings.
|
||||
:type embedding_fn: Callable[[list[str]], list[str]]
|
||||
:raises ValueError: Embedding function is not callable.
|
||||
"""
|
||||
if not hasattr(embedding_fn, "__call__"):
|
||||
raise ValueError("Embedding function is not a function")
|
||||
self.embedding_fn = embedding_fn
|
||||
|
||||
def set_vector_dimension(self, vector_dimension: int):
|
||||
"""
|
||||
Set or overwrite the vector dimension size
|
||||
|
||||
:param vector_dimension: vector dimension size
|
||||
:type vector_dimension: int
|
||||
"""
|
||||
if not isinstance(vector_dimension, int):
|
||||
raise TypeError("vector dimension must be int")
|
||||
self.vector_dimension = vector_dimension
|
||||
|
||||
@staticmethod
|
||||
def _langchain_default_concept(embeddings: Any):
|
||||
"""
|
||||
Langchains default function layout for embeddings.
|
||||
|
||||
:param embeddings: Langchain embeddings
|
||||
:type embeddings: Any
|
||||
:return: embedding function
|
||||
:rtype: Callable
|
||||
"""
|
||||
|
||||
return EmbeddingFunc(embeddings.embed_documents)
|
||||
|
||||
def to_embeddings(self, data: str, **_):
|
||||
"""
|
||||
Convert data to embeddings
|
||||
|
||||
:param data: data to convert to embeddings
|
||||
:type data: str
|
||||
:return: embeddings
|
||||
:rtype: list[float]
|
||||
"""
|
||||
embeddings = self.embedding_fn([data])
|
||||
return embeddings[0]
|
||||
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
|
||||
from chromadb import EmbeddingFunction, Embeddings
|
||||
|
||||
|
||||
class ClarifaiEmbeddingFunction(EmbeddingFunction):
|
||||
def __init__(self, config: BaseEmbedderConfig) -> None:
|
||||
super().__init__()
|
||||
try:
|
||||
from clarifai.client.model import Model
|
||||
from clarifai.client.input import Inputs
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for ClarifaiEmbeddingFunction are not installed."
|
||||
'Please install with `pip install --upgrade "embedchain[clarifai]"`'
|
||||
) from None
|
||||
self.config = config
|
||||
self.api_key = config.api_key or os.getenv("CLARIFAI_PAT")
|
||||
self.model = config.model
|
||||
self.model_obj = Model(url=self.model, pat=self.api_key)
|
||||
self.input_obj = Inputs(pat=self.api_key)
|
||||
|
||||
def __call__(self, input: Union[str, list[str]]) -> Embeddings:
|
||||
if isinstance(input, str):
|
||||
input = [input]
|
||||
|
||||
batch_size = 32
|
||||
embeddings = []
|
||||
try:
|
||||
for i in range(0, len(input), batch_size):
|
||||
batch = input[i : i + batch_size]
|
||||
input_batch = [
|
||||
self.input_obj.get_text_input(input_id=str(id), raw_text=inp) for id, inp in enumerate(batch)
|
||||
]
|
||||
response = self.model_obj.predict(input_batch)
|
||||
embeddings.extend([list(output.data.embeddings[0].vector) for output in response.outputs])
|
||||
except Exception as e:
|
||||
print(f"Predict failed, exception: {e}")
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
class ClarifaiEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config)
|
||||
|
||||
embedding_func = ClarifaiEmbeddingFunction(config=self.config)
|
||||
self.set_embedding_fn(embedding_fn=embedding_func)
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import Optional
|
||||
|
||||
from langchain_cohere.embeddings import CohereEmbeddings
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class CohereEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
embeddings = CohereEmbeddings(model=self.config.model)
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.COHERE.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,38 @@
|
||||
from typing import Optional, Union
|
||||
|
||||
import google.generativeai as genai
|
||||
from chromadb import EmbeddingFunction, Embeddings
|
||||
|
||||
from embedchain.config.embedder.google import GoogleAIEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class GoogleAIEmbeddingFunction(EmbeddingFunction):
|
||||
def __init__(self, config: Optional[GoogleAIEmbedderConfig] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config or GoogleAIEmbedderConfig()
|
||||
|
||||
def __call__(self, input: Union[list[str], str]) -> Embeddings:
|
||||
model = self.config.model
|
||||
title = self.config.title
|
||||
task_type = self.config.task_type
|
||||
if isinstance(input, str):
|
||||
input_ = [input]
|
||||
else:
|
||||
input_ = input
|
||||
data = genai.embed_content(model=model, content=input_, task_type=task_type, title=title)
|
||||
embeddings = data["embedding"]
|
||||
if isinstance(input_, str):
|
||||
embeddings = [embeddings]
|
||||
return embeddings
|
||||
|
||||
|
||||
class GoogleAIEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[GoogleAIEmbedderConfig] = None):
|
||||
super().__init__(config)
|
||||
embedding_fn = GoogleAIEmbeddingFunction(config=config)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.GOOGLE_AI.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,20 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class GPT4AllEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
from langchain_community.embeddings import GPT4AllEmbeddings as LangchainGPT4AllEmbeddings
|
||||
|
||||
model_name = self.config.model or "all-MiniLM-L6-v2-f16.gguf"
|
||||
embeddings = LangchainGPT4AllEmbeddings(model_name=model_name)
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.GPT4ALL.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from langchain_community.embeddings import HuggingFaceEmbeddings
|
||||
|
||||
try:
|
||||
from langchain_huggingface import HuggingFaceEndpointEmbeddings
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for HuggingFaceHub are not installed."
|
||||
"Please install with `pip install langchain_huggingface`"
|
||||
) from None
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class HuggingFaceEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
if self.config.endpoint:
|
||||
if not self.config.api_key and "HUGGINGFACE_ACCESS_TOKEN" not in os.environ:
|
||||
raise ValueError(
|
||||
"Please set the HUGGINGFACE_ACCESS_TOKEN environment variable or pass API Key in the config."
|
||||
)
|
||||
|
||||
embeddings = HuggingFaceEndpointEmbeddings(
|
||||
model=self.config.endpoint,
|
||||
huggingfacehub_api_token=self.config.api_key or os.getenv("HUGGINGFACE_ACCESS_TOKEN"),
|
||||
)
|
||||
else:
|
||||
embeddings = HuggingFaceEmbeddings(model_name=self.config.model, model_kwargs=self.config.model_kwargs)
|
||||
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.HUGGING_FACE.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,46 @@
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
from chromadb import EmbeddingFunction, Embeddings
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class MistralAIEmbeddingFunction(EmbeddingFunction):
|
||||
def __init__(self, config: BaseEmbedderConfig) -> None:
|
||||
super().__init__()
|
||||
try:
|
||||
from langchain_mistralai import MistralAIEmbeddings
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for MistralAI are not installed."
|
||||
'Please install with `pip install --upgrade "embedchain[mistralai]"`'
|
||||
) from None
|
||||
self.config = config
|
||||
api_key = self.config.api_key or os.getenv("MISTRAL_API_KEY")
|
||||
self.client = MistralAIEmbeddings(mistral_api_key=api_key)
|
||||
self.client.model = self.config.model
|
||||
|
||||
def __call__(self, input: Union[list[str], str]) -> Embeddings:
|
||||
if isinstance(input, str):
|
||||
input_ = [input]
|
||||
else:
|
||||
input_ = input
|
||||
response = self.client.embed_documents(input_)
|
||||
return response
|
||||
|
||||
|
||||
class MistralAIEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config)
|
||||
|
||||
if self.config.model is None:
|
||||
self.config.model = "mistral-embed"
|
||||
|
||||
embedding_fn = MistralAIEmbeddingFunction(config=self.config)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.MISTRAL_AI.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,28 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NvidiaEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
if "NVIDIA_API_KEY" not in os.environ:
|
||||
raise ValueError("NVIDIA_API_KEY environment variable must be set")
|
||||
|
||||
super().__init__(config=config)
|
||||
|
||||
model = self.config.model or "nvolveqa_40k"
|
||||
logger.info(f"Using NVIDIA embedding model: {model}")
|
||||
embedder = NVIDIAEmbeddings(model=model)
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embedder)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.NVIDIA_AI.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,32 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
try:
|
||||
from ollama import Client
|
||||
except ImportError:
|
||||
raise ImportError("Ollama Embedder requires extra dependencies. Install with `pip install ollama`") from None
|
||||
|
||||
from langchain_community.embeddings import OllamaEmbeddings
|
||||
|
||||
from embedchain.config import OllamaEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OllamaEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[OllamaEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
client = Client(host=config.base_url)
|
||||
local_models = client.list()["models"]
|
||||
if not any(model.get("name") == self.config.model for model in local_models):
|
||||
logger.info(f"Pulling {self.config.model} from Ollama!")
|
||||
client.pull(self.config.model)
|
||||
embeddings = OllamaEmbeddings(model=self.config.model, base_url=config.base_url)
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.OLLAMA.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
|
||||
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class OpenAIEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
if self.config.model is None:
|
||||
self.config.model = "text-embedding-ada-002"
|
||||
|
||||
api_key = self.config.api_key or os.environ["OPENAI_API_KEY"]
|
||||
api_base = self.config.api_base or os.environ.get("OPENAI_API_BASE")
|
||||
|
||||
if api_key is None and os.getenv("OPENAI_ORGANIZATION") is None:
|
||||
raise ValueError("OPENAI_API_KEY or OPENAI_ORGANIZATION environment variables not provided") # noqa:E501
|
||||
embedding_fn = OpenAIEmbeddingFunction(
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
organization_id=os.getenv("OPENAI_ORGANIZATION"),
|
||||
model_name=self.config.model,
|
||||
)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.OPENAI.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import Optional
|
||||
|
||||
from langchain_community.embeddings import VertexAIEmbeddings
|
||||
|
||||
from embedchain.config import BaseEmbedderConfig
|
||||
from embedchain.embedder.base import BaseEmbedder
|
||||
from embedchain.models import VectorDimensions
|
||||
|
||||
|
||||
class VertexAIEmbedder(BaseEmbedder):
|
||||
def __init__(self, config: Optional[BaseEmbedderConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
embeddings = VertexAIEmbeddings(model_name=config.model)
|
||||
embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
|
||||
self.set_embedding_fn(embedding_fn=embedding_fn)
|
||||
|
||||
vector_dimension = self.config.vector_dimension or VectorDimensions.VERTEX_AI.value
|
||||
self.set_vector_dimension(vector_dimension=vector_dimension)
|
||||
Reference in New Issue
Block a user