biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add Model Revision Support (#1014)

Browse Source 
Co-authored-by: Jasmond Loh <Jasmond.Loh@hotmail.com>
Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>
This commit is contained in:
Jasmond L
2023-09-14 06:20:02 +08:00
committed by GitHub
parent 9841d48a10
commit ab019eea75
20 changed files with 75 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
vllm/transformers_utils/config.py
View File

@@ -1,3 +1,5 @@
from typing import Optional
from transformers import AutoConfig, PretrainedConfig
from vllm.transformers_utils.configs import * # pylint: disable=wildcard-import
@@ -12,10 +14,12 @@ _CONFIG_REGISTRY = {
}
def get_config(model: str, trust_remote_code: bool) -> PretrainedConfig:
def get_config(model: str,
trust_remote_code: bool,
revision: Optional[str] = None) -> PretrainedConfig:
try:
config = AutoConfig.from_pretrained(
model, trust_remote_code=trust_remote_code)
model, trust_remote_code=trust_remote_code, revision=revision)
except ValueError as e:
if (not trust_remote_code and
"requires you to execute the configuration file" in str(e)):
@@ -29,5 +33,5 @@ def get_config(model: str, trust_remote_code: bool) -> PretrainedConfig:
raise e
if config.model_type in _CONFIG_REGISTRY:
config_class = _CONFIG_REGISTRY[config.model_type]
config = config_class.from_pretrained(model)
config = config_class.from_pretrained(model, revision=revision)
return config