diff --git a/src/config/__init__.py b/src/config/__init__.py index 1828e406..217b81f4 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -79,12 +79,12 @@ class Config(SimpleConfig): def handle_self(self): ### handle local model - model_root_dir = os.getenv("MODEL_ROOT_DIR", "pretrained_models") - if self.model_local_paths is not None: - for model, model_rel_path in self.model_local_paths.items(): - # 如果 model_rel_path 不是绝对路径,那么拼接 model_root_dir - if not model_rel_path.startswith("/"): - self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path) + # model_root_dir = os.getenv("MODEL_ROOT_DIR", None) + # if self.model_local_paths is not None: + # for model, model_rel_path in self.model_local_paths.items(): + # # 如果 model_rel_path 不是绝对路径,那么拼接 model_root_dir + # if not model_rel_path.startswith("/"): + # self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path) self.model_names = MODEL_NAMES diff --git a/src/models/README.md b/src/models/README.md index fe1fc347..1594a285 100644 --- a/src/models/README.md +++ b/src/models/README.md @@ -39,16 +39,14 @@ python -m vllm.entrypoints.openai.api_server \ |`zhipu`|`embedding-2`|`ZHIPUAPI` (`.env`)| - - -例如(`config/base.yaml`): +例如(`saves/config/config.yaml`): ```yaml model_provider: qianfan model_name: null # for default -## model dir 可以写相对路径和绝对路径 -### 相对路径是相对于环境变量 (.env) 中 MODEL_ROOT_DIR 的路径 +## model dir 可以写**相对路径**和**绝对路径** +### 相对路径是相对于环境变量 (.env) 中 MODEL_ROOT_DIR(若为空则是相对于 `pretrained_models`) 的路径 model_local_paths: bge-large-zh-v1.5: /models/bge-large-zh-v1.5 ``` diff --git a/src/models/embedding.py b/src/models/embedding.py index 21d32a03..1826ab27 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -15,7 +15,11 @@ GLOBAL_EMBED_STATE = {} class EmbeddingModel(FlagModel): def __init__(self, model_info, config, **kwargs): self.info = model_info - model_name_or_path = config.model_local_paths.get(model_info.name, model_info.default_path) + model_name_or_path = handle_local_model( + paths=config.model_local_paths, + model_name=model_info.name, + default_path=model_info.default_path) + logger.info(f"Loading embedding model {model_info.name} from {model_name_or_path}") super().__init__(model_name_or_path, @@ -30,7 +34,11 @@ class Reranker(FlagReranker): assert config.reranker in RERANKER_LIST.keys(), f"Unsupported Reranker: {config.reranker}, only support {RERANKER_LIST.keys()}" - model_name_or_path = config.model_local_paths.get(config.reranker, RERANKER_LIST[config.reranker]) + model_name_or_path = handle_local_model( + paths=config.model_local_paths, + model_name=config.reranker, + default_path=RERANKER_LIST[config.reranker]) + logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}") super().__init__(model_name_or_path, use_fp16=True, **kwargs) @@ -100,4 +108,10 @@ def get_embedding_model(config): if config.embed_model in ["zhipu-embedding-2", "zhipu-embedding-3"]: model = ZhipuEmbedding(EMBED_MODEL_INFO[config.embed_model], config) - return model \ No newline at end of file + return model + +def handle_local_model(paths, model_name, default_path): + model_path = paths.get(model_name, default_path) + if os.getenv("MODEL_ROOT_DIR") and not os.path.isabs(model_path): + model_path = os.path.join(os.getenv("MODEL_ROOT_DIR"), model_path) + return model_path \ No newline at end of file