model_overwrite

vllm/envs.py

@@ -96,6 +96,7 @@
     VLLM_DP_SIZE: int = 1
     VLLM_DP_MASTER_IP: str = ""
     VLLM_DP_MASTER_PORT: int = 0
+    VLLM_MODEL_OVERWRITE_PATH: Optional[str] = None
 
 
 def get_default_cache_root():
@@ -630,6 +631,11 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
     # Whether to use S3 path for model loading in CI via RunAI Streamer
     "VLLM_CI_USE_S3":
     lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1",
+
+    # Use model_overwrite to redirect the model name to a local folder
+    # while keeping the model name in the test file without being hardcoded.
+    "VLLM_MODEL_OVERWRITE_PATH":
+    lambda: os.environ.get("VLLM_MODEL_OVERWRITE_PATH", None),
 }
 
 # end-env-vars-definition

vllm/transformers_utils/config.py

@@ -24,6 +24,7 @@
     MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
 from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME
 
+from vllm import envs
 from vllm.envs import VLLM_USE_MODELSCOPE
 from vllm.logger import init_logger
 # yapf conflicts with isort for this block
@@ -247,6 +248,8 @@ def get_config(
 ) -> PretrainedConfig:
     # Separate model folder from file path for GGUF models
 
+    model = model_overwrite(model)
+
     is_gguf = check_gguf_file(model)
     if is_gguf:
         kwargs["gguf_file"] = Path(model).name
@@ -758,3 +761,37 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):
         return resolve_obj_by_qualname(function_name)()
     else:
         return nn.Sigmoid() if config.num_labels == 1 else nn.Identity()
+
+
+@cache
+def model_overwrite(model: str):
+    """
+    Use model_overwrite to redirect the model name to a local folder
+    while keeping the model name in the test file without being hardcoded.
+
+    :param model: hf model name
+    :return: maybe overwrite to a local folder
+    """
+
+    model_overwrite_path = envs.VLLM_MODEL_OVERWRITE_PATH
+
+    if not model_overwrite_path:
+        return model
+
+    import pathlib
+
+    if not pathlib.Path(model_overwrite_path.exists()):
+        return model
+
+    model_overwrite_dict = {}
+    with open(model_overwrite_path) as f:
+        for line in f.readlines():
+            model_name, overwrite_name = line.split("\t")
+            model_overwrite_dict[model_name] = overwrite_name.strip()
+
+    if model in model_overwrite_dict:
+        new_model = model_overwrite_dict[model]
+        logger.info("model overwrite: [ %s ] -> [ % ]", model, new_model)
+        model = new_model
+
+    return model

vllm/transformers_utils/tokenizer.py

@@ -15,6 +15,7 @@
 from vllm.envs import VLLM_USE_MODELSCOPE
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
+from vllm.transformers_utils.config import model_overwrite
 from vllm.transformers_utils.tokenizer_base import (TokenizerBase,
                                                     TokenizerRegistry)
 from vllm.transformers_utils.tokenizers import MistralTokenizer
@@ -144,6 +145,9 @@ def get_tokenizer(
 ) -> AnyTokenizer:
     """Gets a tokenizer for the given model name via HuggingFace or ModelScope.
     """
+
+    tokenizer_name = model_overwrite(tokenizer_name)
+
     if VLLM_USE_MODELSCOPE:
         # download model from ModelScope hub,
         # lazy import so that modelscope is not required for normal use.