codekingpro/portable-devtools

1from functools import partial2from typing import Any, Dict, List, Optional, Sequence3 4from langchain_core.callbacks import (5    AsyncCallbackManagerForLLMRun,6    CallbackManagerForLLMRun,7)8from langchain_core.language_models.llms import LLM9from langchain_core.utils import pre_init10 11 12class CTransformers(LLM):13    """C Transformers LLM models.14 15    To use, you should have the ``ctransformers`` python package installed.16    See https://github.com/marella/ctransformers17 18    Example:19        .. code-block:: python20 21            from langchain_community.llms import CTransformers22 23            llm = CTransformers(model="/path/to/ggml-gpt-2.bin", model_type="gpt2")24    """25 26    client: Any  #: :meta private:27 28    model: str29    """The path to a model file or directory or the name of a Hugging Face Hub30    model repo."""31 32    model_type: Optional[str] = None33    """The model type."""34 35    model_file: Optional[str] = None36    """The name of the model file in repo or directory."""37 38    config: Optional[Dict[str, Any]] = None39    """The config parameters.40    See https://github.com/marella/ctransformers#config"""41 42    lib: Optional[str] = None43    """The path to a shared library or one of `avx2`, `avx`, `basic`."""44 45    @property46    def _identifying_params(self) -> Dict[str, Any]:47        """Get the identifying parameters."""48        return {49            "model": self.model,50            "model_type": self.model_type,51            "model_file": self.model_file,52            "config": self.config,53        }54 55    @property56    def _llm_type(self) -> str:57        """Return type of llm."""58        return "ctransformers"59 60    @pre_init61    def validate_environment(cls, values: Dict) -> Dict:62        """Validate that ``ctransformers`` package is installed."""63        try:64            from ctransformers import AutoModelForCausalLM65        except ImportError:66            raise ImportError(67                "Could not import `ctransformers` package. "68                "Please install it with `pip install ctransformers`"69            )70 71        config = values["config"] or {}72        values["client"] = AutoModelForCausalLM.from_pretrained(73            values["model"],74            model_type=values["model_type"],75            model_file=values["model_file"],76            lib=values["lib"],77            **config,78        )79        return values80 81    def _call(82        self,83        prompt: str,84        stop: Optional[Sequence[str]] = None,85        run_manager: Optional[CallbackManagerForLLMRun] = None,86        **kwargs: Any,87    ) -> str:88        """Generate text from a prompt.89 90        Args:91            prompt: The prompt to generate text from.92            stop: A list of sequences to stop generation when encountered.93 94        Returns:95            The generated text.96 97        Example:98            .. code-block:: python99 100                response = llm.invoke("Tell me a joke.")101        """102        text = []103        _run_manager = run_manager or CallbackManagerForLLMRun.get_noop_manager()104        for chunk in self.client(prompt, stop=stop, stream=True):105            text.append(chunk)106            _run_manager.on_llm_new_token(chunk, verbose=self.verbose)107        return "".join(text)108 109    async def _acall(110        self,111        prompt: str,112        stop: Optional[List[str]] = None,113        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,114        **kwargs: Any,115    ) -> str:116        """Asynchronous Call out to CTransformers generate method.117        Very helpful when streaming (like with websockets!)118 119        Args:120            prompt: The prompt to pass into the model.121            stop: A list of strings to stop generation when encountered.122 123        Returns:124            The string generated by the model.125 126        Example:127            .. code-block:: python128                response = llm.invoke("Once upon a time, ")129        """130        text_callback = None131        if run_manager:132            text_callback = partial(run_manager.on_llm_new_token, verbose=self.verbose)133 134        text = ""135        for token in self.client(prompt, stop=stop, stream=True):136            if text_callback:137                await text_callback(token)138            text += token139 140        return text141