mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-12 14:03:48 +08:00
16d2467ea8
* Update TensorRT-LLM --------- Co-authored-by: Denis Kayshev <topenkoff@gmail.com> Co-authored-by: akhoroshev <arthoroshev@gmail.com> Co-authored-by: Patrick Reiter Horn <patrick.horn@gmail.com> Update
32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
from pathlib import Path
|
|
from typing import Any, Literal, Optional, Union
|
|
|
|
from transformers import PreTrainedTokenizerBase
|
|
|
|
from ..llmapi.llm import LLM as BaseLLM
|
|
from ..llmapi.llm import TokenizerBase
|
|
|
|
|
|
class LLM(BaseLLM):
|
|
|
|
def __init__(self,
|
|
model: str,
|
|
tokenizer: Optional[Union[str, Path, TokenizerBase,
|
|
PreTrainedTokenizerBase]] = None,
|
|
tokenizer_mode: Literal['auto', 'slow'] = 'auto',
|
|
skip_tokenizer_init: bool = False,
|
|
trust_remote_code: bool = False,
|
|
tensor_parallel_size: int = 1,
|
|
dtype: str = "auto",
|
|
revision: Optional[str] = None,
|
|
tokenizer_revision: Optional[str] = None,
|
|
speculative_model: Optional[str] = None,
|
|
**kwargs: Any):
|
|
|
|
kwargs_dict = dict(kwargs)
|
|
kwargs_dict['backend'] = 'pytorch'
|
|
super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init,
|
|
trust_remote_code, tensor_parallel_size, dtype,
|
|
revision, tokenizer_revision, speculative_model,
|
|
**kwargs_dict)
|