mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-28 23:07:11 +00:00
844955d6e1
Assigning missed defaults in various classes. Most clients were being
assigned during the `model_validator(mode="before")` step, so this
change should amount to a no-op in those cases.
---
This PR was autogenerated using gritql
```shell
grit apply 'class_definition(name=$C, $body, superclasses=$S) where {
$C <: ! "Config", // Does not work in this scope, but works after class_definition
$body <: block($statements),
$statements <: some bubble assignment(left=$x, right=$y, type=$t) as $A where {
or {
$y <: `Field($z)`,
$x <: "model_config"
}
},
// And has either Any or Optional fields without a default
$statements <: some bubble assignment(left=$x, right=$y, type=$t) as $A where {
$t <: or {
r"Optional.*",
r"Any",
r"Union[None, .*]",
r"Union[.*, None, .*]",
r"Union[.*, None]",
},
$y <: ., // Match empty node
$t => `$t = None`,
},
}
' --language python .
```
130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
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.utils import pre_init
|
|
from pydantic import Field
|
|
|
|
|
|
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 = None #: :meta private:
|
|
|
|
tokenizer: Any = None #: :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"
|