Source code for langchain_community.embeddings.llamacpp

from typing import Any, Dict, List, Optional

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


class LlamaCppEmbeddings(BaseModel, Embeddings):
    """llama.cpp嵌入模型。

    要使用，您应该已安装llama-cpp-python库，并将Llama模型的路径作为构造函数的命名参数提供。
    查看：https://github.com/abetlen/llama-cpp-python

    示例：
        .. code-block:: python

            from langchain_community.embeddings import LlamaCppEmbeddings
            llama = LlamaCppEmbeddings(model_path="/path/to/model.bin")"""

    client: Any  #: :meta private:
    model_path: str

    n_ctx: int = Field(512, alias="n_ctx")
    """标记上下文窗口。"""

    n_parts: int = Field(-1, alias="n_parts")
    """将模型分割成的部分数量。
    如果是-1，则部分数量会自动确定。"""

    seed: int = Field(-1, alias="seed")
    """种子。如果为-1，则使用随机种子。"""

    f16_kv: bool = Field(False, alias="f16_kv")
    """为键/值缓存使用半精度。"""

    logits_all: bool = Field(False, alias="logits_all")
    """返回所有标记的logits，而不仅仅是最后一个标记。"""

    vocab_only: bool = Field(False, alias="vocab_only")
    """只加载词汇表，不加载权重。"""

    use_mlock: bool = Field(False, alias="use_mlock")
    """强制系统将模型保留在内存中。"""

    n_threads: Optional[int] = Field(None, alias="n_threads")
    """线程数。如果为None，则线程数会自动确定。"""

    n_batch: Optional[int] = Field(512, alias="n_batch")
    """并行处理的令牌数量。
应该是介于1和n_ctx之间的数字。"""

    n_gpu_layers: Optional[int] = Field(None, alias="n_gpu_layers")
    """要加载到GPU内存中的层数。默认值为None。"""

    verbose: bool = Field(True, alias="verbose")
    """将详细输出打印到stderr。"""

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证llama-cpp-python库是否已安装。"""
        model_path = values["model_path"]
        model_param_names = [
            "n_ctx",
            "n_parts",
            "seed",
            "f16_kv",
            "logits_all",
            "vocab_only",
            "use_mlock",
            "n_threads",
            "n_batch",
            "verbose",
        ]
        model_params = {k: values[k] for k in model_param_names}
        # For backwards compatibility, only include if non-null.
        if values["n_gpu_layers"] is not None:
            model_params["n_gpu_layers"] = values["n_gpu_layers"]

        try:
            from llama_cpp import Llama

            values["client"] = Llama(model_path, embedding=True, **model_params)
        except ImportError:
            raise ImportError(
                "Could not import llama-cpp-python library. "
                "Please install the llama-cpp-python library to "
                "use this embedding model: pip install llama-cpp-python"
            )
        except Exception as e:
            raise ValueError(
                f"Could not load Llama model from path: {model_path}. "
                f"Received error {e}"
            )

        return values

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

参数:
    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]:
        """使用Llama模型嵌入一个查询。

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

返回：
    文本的嵌入结果。
"""
        embedding = self.client.embed(text)
        return list(map(float, embedding))