refactor: RAG Refactor (#985)

Co-authored-by: Aralhi <xiaoping0501@gmail.com>
Co-authored-by: csunny <cfqsunny@163.com>

dbgpt/rag/embedding/embedding_factory.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Type, TYPE_CHECKING
+
+from dbgpt.component import BaseComponent
+from dbgpt.rag.embedding.embeddings import HuggingFaceEmbeddings
+
+if TYPE_CHECKING:
+    from dbgpt.rag.embedding.embeddings import Embeddings
+
+
+class EmbeddingFactory(BaseComponent, ABC):
+    """Abstract base class for EmbeddingFactory."""
+
+    name = "embedding_factory"
+
+    @abstractmethod
+    def create(
+        self, model_name: str = None, embedding_cls: Type = None
+    ) -> "Embeddings":
+        """Create embedding"""
+
+
+class DefaultEmbeddingFactory(EmbeddingFactory):
+    def __init__(
+        self, system_app=None, default_model_name: str = None, **kwargs: Any
+    ) -> None:
+        super().__init__(system_app=system_app)
+        self._default_model_name = default_model_name
+        self.kwargs = kwargs
+
+    def init_app(self, system_app):
+        pass
+
+    def create(
+        self, model_name: str = None, embedding_cls: Type = None
+    ) -> "Embeddings":
+        if not model_name:
+            model_name = self._default_model_name
+
+        new_kwargs = {k: v for k, v in self.kwargs.items()}
+        new_kwargs["model_name"] = model_name
+
+        if embedding_cls:
+            return embedding_cls(**new_kwargs)
+        else:
+            return HuggingFaceEmbeddings(**new_kwargs)

dbgpt/rag/embedding/embeddings.py

@@ -0,0 +1,363 @@
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+import requests
+from pydantic import Field, Extra, BaseModel
+
+DEFAULT_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
+DEFAULT_INSTRUCT_MODEL = "hkunlp/instructor-large"
+DEFAULT_BGE_MODEL = "BAAI/bge-large-en"
+DEFAULT_EMBED_INSTRUCTION = "Represent the document for retrieval: "
+DEFAULT_QUERY_INSTRUCTION = (
+    "Represent the question for retrieving supporting documents: "
+)
+DEFAULT_QUERY_BGE_INSTRUCTION_EN = (
+    "Represent this question for searching relevant passages: "
+)
+DEFAULT_QUERY_BGE_INSTRUCTION_ZH = "为这个句子生成表示以用于检索相关文章："
+
+
+class Embeddings(ABC):
+    """Interface for embedding models."""
+
+    """refer to https://github.com/langchain-ai/langchain/tree/master/libs/langchain/langchain/embeddings"""
+
+    @abstractmethod
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed search docs."""
+
+    @abstractmethod
+    def embed_query(self, text: str) -> List[float]:
+        """Embed query text."""
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Asynchronous Embed search docs."""
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.embed_documents, texts
+        )
+
+    async def aembed_query(self, text: str) -> List[float]:
+        """Asynchronous Embed query text."""
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.embed_query, text
+        )
+
+
+class HuggingFaceEmbeddings(BaseModel, Embeddings):
+    """HuggingFace sentence_transformers embedding models.
+    To use, you should have the ``sentence_transformers`` python package installed.
+    Refer to https://github.com/langchain-ai/langchain/tree/master/libs/langchain/langchain/embeddings
+    Example:
+        .. code-block:: python
+
+            from .embeddings import HuggingFaceEmbeddings
+
+            model_name = "sentence-transformers/all-mpnet-base-v2"
+            model_kwargs = {'device': 'cpu'}
+            encode_kwargs = {'normalize_embeddings': False}
+            hf = HuggingFaceEmbeddings(
+                model_name=model_name,
+                model_kwargs=model_kwargs,
+                encode_kwargs=encode_kwargs
+            )
+    """
+
+    client: Any  #: :meta private:
+    model_name: str = DEFAULT_MODEL_NAME
+    """Model name to use."""
+    cache_folder: Optional[str] = None
+    """Path to store models. 
+    Can be also set by SENTENCE_TRANSFORMERS_HOME environment variable."""
+    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass to the model."""
+    encode_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass when calling the `encode` method of the model."""
+    multi_process: bool = False
+    """Run encode() on multiple GPUs."""
+
+    def __init__(self, **kwargs: Any):
+        """Initialize the sentence_transformer."""
+        super().__init__(**kwargs)
+        try:
+            import sentence_transformers
+
+        except ImportError as exc:
+            raise ImportError(
+                "Could not import sentence_transformers python package. "
+                "Please install it with `pip install sentence-transformers`."
+            ) from exc
+
+        self.client = sentence_transformers.SentenceTransformer(
+            self.model_name, cache_folder=self.cache_folder, **self.model_kwargs
+        )
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Compute doc embeddings using a HuggingFace transformer model.
+
+        Args:
+            texts: The list of texts to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        import sentence_transformers
+
+        texts = list(map(lambda x: x.replace("\n", " "), texts))
+        if self.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, **self.encode_kwargs)
+
+        return embeddings.tolist()
+
+    def embed_query(self, text: str) -> List[float]:
+        """Compute query embeddings using a HuggingFace transformer model.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+        return self.embed_documents([text])[0]
+
+
+class HuggingFaceInstructEmbeddings(BaseModel, Embeddings):
+    """Wrapper around sentence_transformers embedding models.
+
+    To use, you should have the ``sentence_transformers``
+    and ``InstructorEmbedding`` python packages installed.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.embeddings import HuggingFaceInstructEmbeddings
+
+            model_name = "hkunlp/instructor-large"
+            model_kwargs = {'device': 'cpu'}
+            encode_kwargs = {'normalize_embeddings': True}
+            hf = HuggingFaceInstructEmbeddings(
+                model_name=model_name,
+                model_kwargs=model_kwargs,
+                encode_kwargs=encode_kwargs
+            )
+    """
+
+    client: Any  #: :meta private:
+    model_name: str = DEFAULT_INSTRUCT_MODEL
+    """Model name to use."""
+    cache_folder: Optional[str] = None
+    """Path to store models. 
+    Can be also set by SENTENCE_TRANSFORMERS_HOME environment variable."""
+    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass to the model."""
+    encode_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass when calling the `encode` method of the model."""
+    embed_instruction: str = DEFAULT_EMBED_INSTRUCTION
+    """Instruction to use for embedding documents."""
+    query_instruction: str = DEFAULT_QUERY_INSTRUCTION
+    """Instruction to use for embedding query."""
+
+    def __init__(self, **kwargs: Any):
+        """Initialize the sentence_transformer."""
+        super().__init__(**kwargs)
+        try:
+            from InstructorEmbedding import INSTRUCTOR
+
+            self.client = INSTRUCTOR(
+                self.model_name, cache_folder=self.cache_folder, **self.model_kwargs
+            )
+        except ImportError as e:
+            raise ImportError("Dependencies for InstructorEmbedding not found.") from e
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Compute doc embeddings using a HuggingFace instruct model.
+
+        Args:
+            texts: The list of texts to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        instruction_pairs = [[self.embed_instruction, text] for text in texts]
+        embeddings = self.client.encode(instruction_pairs, **self.encode_kwargs)
+        return embeddings.tolist()
+
+    def embed_query(self, text: str) -> List[float]:
+        """Compute query embeddings using a HuggingFace instruct model.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+        instruction_pair = [self.query_instruction, text]
+        embedding = self.client.encode([instruction_pair], **self.encode_kwargs)[0]
+        return embedding.tolist()
+
+
+class HuggingFaceBgeEmbeddings(BaseModel, Embeddings):
+    """HuggingFace BGE sentence_transformers embedding models.
+
+    To use, you should have the ``sentence_transformers`` python package installed.
+    refer to https://github.com/langchain-ai/langchain/tree/master/libs/langchain/langchain/embeddings
+    Example:
+        .. code-block:: python
+
+            from langchain.embeddings import HuggingFaceBgeEmbeddings
+
+            model_name = "BAAI/bge-large-en"
+            model_kwargs = {'device': 'cpu'}
+            encode_kwargs = {'normalize_embeddings': True}
+            hf = HuggingFaceBgeEmbeddings(
+                model_name=model_name,
+                model_kwargs=model_kwargs,
+                encode_kwargs=encode_kwargs
+            )
+    """
+
+    client: Any  #: :meta private:
+    model_name: str = DEFAULT_BGE_MODEL
+    """Model name to use."""
+    cache_folder: Optional[str] = None
+    """Path to store models.
+    Can be also set by SENTENCE_TRANSFORMERS_HOME environment variable."""
+    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass to the model."""
+    encode_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass when calling the `encode` method of the model."""
+    query_instruction: str = DEFAULT_QUERY_BGE_INSTRUCTION_EN
+    """Instruction to use for embedding query."""
+
+    def __init__(self, **kwargs: Any):
+        """Initialize the sentence_transformer."""
+        super().__init__(**kwargs)
+        try:
+            import sentence_transformers
+
+        except ImportError as exc:
+            raise ImportError(
+                "Could not import sentence_transformers python package. "
+                "Please install it with `pip install sentence_transformers`."
+            ) from exc
+
+        self.client = sentence_transformers.SentenceTransformer(
+            self.model_name, cache_folder=self.cache_folder, **self.model_kwargs
+        )
+        if "-zh" in self.model_name:
+            self.query_instruction = DEFAULT_QUERY_BGE_INSTRUCTION_ZH
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Compute doc embeddings using a HuggingFace transformer model.
+
+        Args:
+            texts: The list of texts to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        texts = [t.replace("\n", " ") for t in texts]
+        embeddings = self.client.encode(texts, **self.encode_kwargs)
+        return embeddings.tolist()
+
+    def embed_query(self, text: str) -> List[float]:
+        """Compute query embeddings using a HuggingFace transformer model.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+        text = text.replace("\n", " ")
+        embedding = self.client.encode(
+            self.query_instruction + text, **self.encode_kwargs
+        )
+        return embedding.tolist()
+
+
+class HuggingFaceInferenceAPIEmbeddings(BaseModel, Embeddings):
+    """Embed texts using the HuggingFace API.
+
+    Requires a HuggingFace Inference API key and a model name.
+    """
+
+    api_key: str
+    """Your API key for the HuggingFace Inference API."""
+    model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
+    """The name of the model to use for text embeddings."""
+
+    @property
+    def _api_url(self) -> str:
+        return (
+            "https://api-inference.huggingface.co"
+            "/pipeline"
+            "/feature-extraction"
+            f"/{self.model_name}"
+        )
+
+    @property
+    def _headers(self) -> dict:
+        return {"Authorization": f"Bearer {self.api_key}"}
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Get the embeddings for a list of texts.
+
+        Args:
+            texts (Documents): A list of texts to get embeddings for.
+
+        Returns:
+            Embedded texts as List[List[float]], where each inner List[float]
+                corresponds to a single input text.
+
+        Example:
+            .. code-block:: python
+
+                from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
+
+                hf_embeddings = HuggingFaceInferenceAPIEmbeddings(
+                    api_key="your_api_key",
+                    model_name="sentence-transformers/all-MiniLM-l6-v2"
+                )
+                texts = ["Hello, world!", "How are you?"]
+                hf_embeddings.embed_documents(texts)
+        """
+        response = requests.post(
+            self._api_url,
+            headers=self._headers,
+            json={
+                "inputs": texts,
+                "options": {"wait_for_model": True, "use_cache": True},
+            },
+        )
+        return response.json()
+
+    def embed_query(self, text: str) -> List[float]:
+        """Compute query embeddings using a HuggingFace transformer model.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+        return self.embed_documents([text])[0]