Source code for langchain_community.embeddings.gpt4all

from typing import Any, Dict, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, root_validator


class GPT4AllEmbeddings(BaseModel, Embeddings):
    """GPT4All嵌入模型。

    要使用，您应该安装gpt4all python包

    示例:
        .. code-block:: python

            from langchain_community.embeddings import GPT4AllEmbeddings

            model_name = "all-MiniLM-L6-v2.gguf2.f16.gguf"
            gpt4all_kwargs = {'allow_download': 'True'}
            embeddings = GPT4AllEmbeddings(
                model_name=model_name,
                gpt4all_kwargs=gpt4all_kwargs
            )
"""

    model_name: str
    n_threads: Optional[int] = None
    device: Optional[str] = "cpu"
    gpt4all_kwargs: Optional[dict] = {}
    client: Any  #: :meta private:

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证是否安装了GPT4All库。"""

        try:
            from gpt4all import Embed4All

            values["client"] = Embed4All(
                model_name=values["model_name"],
                n_threads=values.get("n_threads"),
                device=values.get("device"),
                **values.get("gpt4all_kwargs"),
            )
        except ImportError:
            raise ImportError(
                "Could not import gpt4all library. "
                "Please install the gpt4all library to "
                "use this embedding model: pip install gpt4all"
            )
        return values

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """使用GPT4All嵌入文档列表。

参数：
    texts：要嵌入的文本列表。

返回：
    每个文本的嵌入列表。
"""

        embeddings = [self.client.embed(text) for text in texts]
        return [list(map(float, e)) for e in embeddings]

    def embed_query(self, text: str) -> List[float]:
        """使用GPT4All嵌入一个查询。

参数：
    text: 要嵌入的文本。

返回：
    文本的嵌入结果。
"""
        return self.embed_documents([text])[0]