Source code for langchain_community.llms.ctranslate2

from typing import Any, Dict, List, Optional, Union

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import BaseLLM
from langchain_core.outputs import Generation, LLMResult
from langchain_core.pydantic_v1 import Field
from langchain_core.utils import pre_init


[docs]class CTranslate2(BaseLLM): """CTranslate2 language model.""" model_path: str = "" """Path to the CTranslate2 model directory.""" tokenizer_name: str = "" """Name of the original Hugging Face model needed to load the proper tokenizer.""" device: str = "cpu" """Device to use (possible values are: cpu, cuda, auto).""" device_index: Union[int, List[int]] = 0 """Device IDs where to place this generator on.""" compute_type: Union[str, Dict[str, str]] = "default" """ Model computation type or a dictionary mapping a device name to the computation type (possible values are: default, auto, int8, int8_float32, int8_float16, int8_bfloat16, int16, float16, bfloat16, float32). """ max_length: int = 512 """Maximum generation length.""" sampling_topk: int = 1 """Randomly sample predictions from the top K candidates.""" sampling_topp: float = 1 """Keep the most probable tokens whose cumulative probability exceeds this value.""" sampling_temperature: float = 1 """Sampling temperature to generate more random samples.""" client: Any #: :meta private: tokenizer: Any #: :meta private: ctranslate2_kwargs: Dict[str, Any] = Field(default_factory=dict) """ Holds any model parameters valid for `ctranslate2.Generator` call not explicitly specified. """ @pre_init def validate_environment(cls, values: Dict) -> Dict: """Validate that python package exists in environment.""" try: import ctranslate2 except ImportError: raise ImportError( "Could not import ctranslate2 python package. " "Please install it with `pip install ctranslate2`." ) try: import transformers except ImportError: raise ImportError( "Could not import transformers python package. " "Please install it with `pip install transformers`." ) values["client"] = ctranslate2.Generator( model_path=values["model_path"], device=values["device"], device_index=values["device_index"], compute_type=values["compute_type"], **values["ctranslate2_kwargs"], ) values["tokenizer"] = transformers.AutoTokenizer.from_pretrained( values["tokenizer_name"] ) return values @property def _default_params(self) -> Dict[str, Any]: """Get the default parameters.""" return { "max_length": self.max_length, "sampling_topk": self.sampling_topk, "sampling_topp": self.sampling_topp, "sampling_temperature": self.sampling_temperature, } def _generate( self, prompts: List[str], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any, ) -> LLMResult: # build sampling parameters params = {**self._default_params, **kwargs} # call the model encoded_prompts = self.tokenizer(prompts)["input_ids"] tokenized_prompts = [ self.tokenizer.convert_ids_to_tokens(encoded_prompt) for encoded_prompt in encoded_prompts ] results = self.client.generate_batch(tokenized_prompts, **params) sequences = [result.sequences_ids[0] for result in results] decoded_sequences = [self.tokenizer.decode(seq) for seq in sequences] generations = [] for text in decoded_sequences: generations.append([Generation(text=text)]) return LLMResult(generations=generations) @property def _llm_type(self) -> str: """Return type of llm.""" return "ctranslate2"