codekingpro/portable-devtools

1from typing import Any, List, Optional2 3from langchain_core.embeddings import Embeddings4from pydantic import BaseModel, ConfigDict, Field, model_validator5from typing_extensions import Self6 7 8class LlamaCppEmbeddings(BaseModel, Embeddings):9    """llama.cpp embedding models.10 11    To use, you should have the llama-cpp-python library installed, and provide the12    path to the Llama model as a named parameter to the constructor.13    Check out: https://github.com/abetlen/llama-cpp-python14 15    Example:16        .. code-block:: python17 18            from langchain_community.embeddings import LlamaCppEmbeddings19            llama = LlamaCppEmbeddings(model_path="/path/to/model.bin")20    """21 22    client: Any = None  #: :meta private:23    model_path: str = Field(default="")24 25    n_ctx: int = Field(512, alias="n_ctx")26    """Token context window."""27 28    n_parts: int = Field(-1, alias="n_parts")29    """Number of parts to split the model into. 30    If -1, the number of parts is automatically determined."""31 32    seed: int = Field(-1, alias="seed")33    """Seed. If -1, a random seed is used."""34 35    f16_kv: bool = Field(False, alias="f16_kv")36    """Use half-precision for key/value cache."""37 38    logits_all: bool = Field(False, alias="logits_all")39    """Return logits for all tokens, not just the last token."""40 41    vocab_only: bool = Field(False, alias="vocab_only")42    """Only load the vocabulary, no weights."""43 44    use_mlock: bool = Field(False, alias="use_mlock")45    """Force system to keep model in RAM."""46 47    n_threads: Optional[int] = Field(None, alias="n_threads")48    """Number of threads to use. If None, the number 49    of threads is automatically determined."""50 51    n_batch: Optional[int] = Field(512, alias="n_batch")52    """Number of tokens to process in parallel.53    Should be a number between 1 and n_ctx."""54 55    n_gpu_layers: Optional[int] = Field(None, alias="n_gpu_layers")56    """Number of layers to be loaded into gpu memory. Default None."""57 58    verbose: bool = Field(True, alias="verbose")59    """Print verbose output to stderr."""60 61    device: Optional[str] = Field(None, alias="device")62    """Device type to use and pass to the model"""63 64    model_config = ConfigDict(65        extra="forbid",66        protected_namespaces=(),67    )68 69    @model_validator(mode="after")70    def validate_environment(self) -> Self:71        """Validate that llama-cpp-python library is installed."""72        model_path = self.model_path73        model_param_names = [74            "n_ctx",75            "n_parts",76            "seed",77            "f16_kv",78            "logits_all",79            "vocab_only",80            "use_mlock",81            "n_threads",82            "n_batch",83            "verbose",84            "device",85        ]86        model_params = {k: getattr(self, k) for k in model_param_names}87        # For backwards compatibility, only include if non-null.88        if self.n_gpu_layers is not None:89            model_params["n_gpu_layers"] = self.n_gpu_layers90 91        if not self.client:92            try:93                from llama_cpp import Llama94 95                self.client = Llama(model_path, embedding=True, **model_params)96            except ImportError:97                raise ImportError(98                    "Could not import llama-cpp-python library. "99                    "Please install the llama-cpp-python library to "100                    "use this embedding model: pip install llama-cpp-python"101                )102            except Exception as e:103                raise ValueError(104                    f"Could not load Llama model from path: {model_path}. "105                    f"Received error {e}"106                )107 108        return self109 110    def embed_documents(self, texts: List[str]) -> List[List[float]]:111        """Embed a list of documents using the Llama model.112 113        Args:114            texts: The list of texts to embed.115 116        Returns:117            List of embeddings, one for each text.118        """119        embeddings = self.client.create_embedding(texts)120        final_embeddings = []121        for e in embeddings["data"]:122            try:123                if isinstance(e["embedding"][0], list):124                    for data in e["embedding"]:125                        final_embeddings.append(list(map(float, data)))126                else:127                    final_embeddings.append(list(map(float, e["embedding"])))128            except (IndexError, TypeError):129                final_embeddings.append(list(map(float, e["embedding"])))130        return final_embeddings131 132    def embed_query(self, text: str) -> List[float]:133        """Embed a query using the Llama model.134 135        Args:136            text: The text to embed.137 138        Returns:139            Embeddings for the text.140        """141        embedding = self.client.embed(text)142        if embedding and isinstance(embedding, list) and isinstance(embedding[0], list):143            return list(map(float, embedding[0]))144        else:145            return list(map(float, embedding))146