from typing import Any, Dict, List, Optional
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseLLM
from langchain_core.outputs import Generation, LLMResult
from langchain_core.utils import pre_init
from pydantic import Field
[docs]
class Aphrodite(BaseLLM):
"""Aphrodite language model."""
model: str = ""
"""The name or path of a HuggingFace Transformers model."""
tensor_parallel_size: Optional[int] = 1
"""The number of GPUs to use for distributed execution with tensor parallelism."""
trust_remote_code: Optional[bool] = False
"""Trust remote code (e.g., from HuggingFace) when downloading the model
and tokenizer."""
n: int = 1
"""Number of output sequences to return for the given prompt."""
best_of: Optional[int] = None
"""Number of output sequences that are generated from the prompt.
From these `best_of` sequences, the top `n` sequences are returned.
`best_of` must be >= `n`. This is treated as the beam width when
`use_beam_search` is True. By default, `best_of` is set to `n`."""
presence_penalty: float = 0.0
"""Float that penalizes new tokens based on whether they appear in the
generated text so far. Values > 0 encourage the model to generate new
tokens, while values < 0 encourage the model to repeat tokens."""
frequency_penalty: float = 0.0
"""Float that penalizes new tokens based on their frequency in the
generated text so far. Applied additively to the logits."""
repetition_penalty: float = 1.0
"""Float that penalizes new tokens based on their frequency in the
generated text so far. Applied multiplicatively to the logits."""
temperature: float = 1.0
"""Float that controls the randomness of the sampling. Lower values
make the model more deterministic, while higher values make the model
more random. Zero is equivalent to greedy sampling."""
top_p: float = 1.0
"""Float that controls the cumulative probability of the top tokens to consider.
Must be in (0, 1]. Set to 1.0 to consider all tokens."""
top_k: int = -1
"""Integer that controls the number of top tokens to consider. Set to -1 to
consider all tokens (disabled)."""
top_a: float = 0.0
"""Float that controls the cutoff for Top-A sampling. Exact cutoff is
top_a*max_prob**2. Must be in [0,inf], 0 to disable."""
min_p: float = 0.0
"""Float that controls the cutoff for min-p sampling. Exact cutoff is
min_p*max_prob. Must be in [0,1], 0 to disable."""
tfs: float = 1.0
"""Float that controls the cumulative approximate curvature of the
distribution to retain for Tail Free Sampling. Must be in (0, 1].
Set to 1.0 to disable."""
eta_cutoff: float = 0.0
"""Float that controls the cutoff threshold for Eta sampling
(a form of entropy adaptive truncation sampling). Threshold is
calculated as `min(eta, sqrt(eta)*entropy(probs)). Specified
in units of 1e-4. Set to 0 to disable."""
epsilon_cutoff: float = 0.0
"""Float that controls the cutoff threshold for Epsilon sampling
(simple probability threshold truncation). Specified in units of
1e-4. Set to 0 to disable."""
typical_p: float = 1.0
"""Float that controls the cumulative probability of tokens closest
in surprise to the expected surprise to consider. Must be in (0, 1].
Set to 1 to disable."""
mirostat_mode: int = 0
"""The mirostat mode to use. 0 for no mirostat, 2 for mirostat v2.
Mode 1 is not supported."""
mirostat_tau: float = 0.0
"""The target 'surprisal' that mirostat works towards. Range [0, inf)."""
use_beam_search: bool = False
"""Whether to use beam search instead of sampling."""
length_penalty: float = 1.0
"""Float that penalizes sequences based on their length. Used only
when `use_beam_search` is True."""
early_stopping: bool = False
"""Controls the stopping condition for beam search. It accepts the
following values: `True`, where the generation stops as soon as there
are `best_of` complete candidates; `False`, where a heuristic is applied
to the generation stops when it is very unlikely to find better candidates;
`never`, where the beam search procedure only stops where there cannot be
better candidates (canonical beam search algorithm)."""
stop: Optional[List[str]] = None
"""List of strings that stop the generation when they are generated.
The returned output will not contain the stop tokens."""
stop_token_ids: Optional[List[int]] = None
"""List of tokens that stop the generation when they are generated.
The returned output will contain the stop tokens unless the stop tokens
are special tokens."""
ignore_eos: bool = False
"""Whether to ignore the EOS token and continue generating tokens after
the EOS token is generated."""
max_tokens: int = 512
"""Maximum number of tokens to generate per output sequence."""
logprobs: Optional[int] = None
"""Number of log probabilities to return per output token."""
prompt_logprobs: Optional[int] = None
"""Number of log probabilities to return per prompt token."""
custom_token_bans: Optional[List[int]] = None
"""List of token IDs to ban from generating."""
skip_special_tokens: bool = True
"""Whether to skip special tokens in the output. Defaults to True."""
spaces_between_special_tokens: bool = True
"""Whether to add spaces between special tokens in the output.
Defaults to True."""
logit_bias: Optional[Dict[str, float]] = None
"""List of LogitsProcessors to change the probability of token
prediction at runtime."""
dtype: str = "auto"
"""The data type for the model weights and activations."""
download_dir: Optional[str] = None
"""Directory to download and load the weights. (Default to the default
cache dir of huggingface)"""
quantization: Optional[str] = None
"""Quantization mode to use. Can be one of `awq` or `gptq`."""
aphrodite_kwargs: Dict[str, Any] = Field(default_factory=dict)
"""Holds any model parameters valid for `aphrodite.LLM` call not explicitly
specified."""
client: Any = None #: :meta private:
[docs]
@pre_init
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that python package exists in environment."""
try:
from aphrodite import LLM as AphroditeModel
except ImportError:
raise ImportError(
"Could not import aphrodite-engine python package. "
"Please install it with `pip install aphrodite-engine`."
)
# aphrodite_kwargs = values["aphrodite_kwargs"]
# if values.get("quantization"):
# aphrodite_kwargs["quantization"] = values["quantization"]
values["client"] = AphroditeModel(
model=values["model"],
tensor_parallel_size=values["tensor_parallel_size"],
trust_remote_code=values["trust_remote_code"],
dtype=values["dtype"],
download_dir=values["download_dir"],
**values["aphrodite_kwargs"],
)
return values
@property
def _default_params(self) -> Dict[str, Any]:
"""Get the default parameters for calling aphrodite."""
return {
"n": self.n,
"best_of": self.best_of,
"max_tokens": self.max_tokens,
"top_k": self.top_k,
"top_p": self.top_p,
"top_a": self.top_a,
"min_p": self.min_p,
"temperature": self.temperature,
"presence_penalty": self.presence_penalty,
"frequency_penalty": self.frequency_penalty,
"repetition_penalty": self.repetition_penalty,
"tfs": self.tfs,
"eta_cutoff": self.eta_cutoff,
"epsilon_cutoff": self.epsilon_cutoff,
"typical_p": self.typical_p,
"mirostat_mode": self.mirostat_mode,
"mirostat_tau": self.mirostat_tau,
"length_penalty": self.length_penalty,
"early_stopping": self.early_stopping,
"use_beam_search": self.use_beam_search,
"stop": self.stop,
"ignore_eos": self.ignore_eos,
"logprobs": self.logprobs,
"prompt_logprobs": self.prompt_logprobs,
"custom_token_bans": self.custom_token_bans,
"skip_special_tokens": self.skip_special_tokens,
"spaces_between_special_tokens": self.spaces_between_special_tokens,
"logit_bias": self.logit_bias,
}
def _generate(
self,
prompts: List[str],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> LLMResult:
"""Run the LLM on the given prompt and input."""
from aphrodite import SamplingParams
# build sampling parameters
params = {**self._default_params, **kwargs, "stop": stop}
if "logit_bias" in params:
del params["logit_bias"]
sampling_params = SamplingParams(**params)
# call the model
outputs = self.client.generate(prompts, sampling_params)
generations = []
for output in outputs:
text = output.outputs[0].text
generations.append([Generation(text=text)])
return LLMResult(generations=generations)
@property
def _llm_type(self) -> str:
"""Return type of llm."""
return "aphrodite"
