mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-14 06:53:50 +08:00
258c7540c0
Co-authored-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> open source f8c0381a2bc50ee2739c3d8c2be481b31e5f00bd (#2736) Co-authored-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Add note for blackwell (#2742) Update the docs to workaround the extra-index-url issue (#2744) update README.md (#2751) Fix github io pages (#2761) 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)
|