[Misc] Improve model redirect to accept json dictionary (#16119)
Signed-off-by: Isotr0py <2037008807@qq.com>
This commit is contained in:
parent
d5ae4f7f42
commit
c2a9671510
@ -665,6 +665,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
|
||||
lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1",
|
||||
|
||||
# Use model_redirect to redirect the model name to a local folder.
|
||||
# `model_redirect` can be a json file mapping the model between
|
||||
# repo_id and local folder:
|
||||
# {"meta-llama/Llama-3.2-1B": "/tmp/Llama-3.2-1B"}
|
||||
# or a space separated values table file:
|
||||
# meta-llama/Llama-3.2-1B /tmp/Llama-3.2-1B
|
||||
"VLLM_MODEL_REDIRECT_PATH":
|
||||
lambda: os.environ.get("VLLM_MODEL_REDIRECT_PATH", None),
|
||||
|
||||
|
@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from functools import cache
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
@ -51,6 +52,26 @@ def modelscope_list_repo_files(
|
||||
return files
|
||||
|
||||
|
||||
def _maybe_json_dict(path: Union[str, PathLike]) -> dict[str, str]:
|
||||
with open(path) as f:
|
||||
try:
|
||||
return json.loads(f.read())
|
||||
except Exception:
|
||||
return dict[str, str]()
|
||||
|
||||
|
||||
def _maybe_space_split_dict(path: Union[str, PathLike]) -> dict[str, str]:
|
||||
parsed_dict = dict[str, str]()
|
||||
with open(path) as f:
|
||||
for line in f.readlines():
|
||||
try:
|
||||
model_name, redirect_name = line.strip().split()
|
||||
parsed_dict[model_name] = redirect_name
|
||||
except Exception:
|
||||
pass
|
||||
return parsed_dict
|
||||
|
||||
|
||||
@cache
|
||||
def maybe_model_redirect(model: str) -> str:
|
||||
"""
|
||||
@ -68,16 +89,10 @@ def maybe_model_redirect(model: str) -> str:
|
||||
if not Path(model_redirect_path).exists():
|
||||
return model
|
||||
|
||||
with open(model_redirect_path) as f:
|
||||
for line in f.readlines():
|
||||
try:
|
||||
model_name, redirect_name = line.split("\t")
|
||||
if model == model_name:
|
||||
redirect_name = redirect_name.strip()
|
||||
logger.info("model redirect: [ %s ] -> [ %s ]", model,
|
||||
redirect_name)
|
||||
return redirect_name
|
||||
except Exception:
|
||||
pass
|
||||
redirect_dict = (_maybe_json_dict(model_redirect_path)
|
||||
or _maybe_space_split_dict(model_redirect_path))
|
||||
if (redirect_model := redirect_dict.get(model)):
|
||||
logger.info("model redirect: [ %s ] -> [ %s ]", model, redirect_model)
|
||||
return redirect_model
|
||||
|
||||
return model
|
||||
|
Loading…
x
Reference in New Issue
Block a user